#[path = "../common/provider.rs"]
mod provider;
use adk_agent::LlmAgentBuilder;
use adk_core::{
Agent, Content, Event, FunctionResponseData, Llm, MultiAgentLoader, Part, SessionId, Tool,
ToolContext, UserId,
};
use adk_runner::{Runner, RunnerConfig};
use adk_server::{ServerConfig, create_app, shutdown_signal};
use adk_session::{CreateRequest, GetRequest, InMemorySessionService, SessionService};
use adk_ui::{
MCP_APPS_PROTOCOL_VERSION, McpAppsAppCapabilities, McpAppsInitializeRequest,
McpAppsInitializeRequestParams, McpAppsInitializeResult, McpAppsPartyInfo,
SUPPORTED_UI_PROTOCOLS, UI_CAPABILITY_SCHEMA_VERSION, UI_DEFAULT_PROTOCOL,
UI_PREFERRED_PROTOCOL, UI_PROTOCOL_CAPABILITIES, UiToolset, a2ui::A2UI_AGENT_PROMPT,
build_default_mcp_apps_initialize_result, default_a2ui_agent_capabilities,
default_a2ui_renderer_capabilities, default_ag_ui_agent_capabilities,
default_mcp_apps_extension_capabilities, default_mcp_apps_host_capabilities,
normalize_runtime_ui_protocol,
};
use anyhow::Result;
use async_trait::async_trait;
use axum::{
body::{Body, to_bytes},
extract::State,
http::{HeaderMap, HeaderValue, Method, Request, StatusCode, header},
middleware::{self, Next},
response::{
IntoResponse, Response,
sse::{Event as SseEvent, KeepAlive, Sse},
},
};
use base64::{Engine as _, engine::general_purpose::STANDARD as BASE64_STANDARD};
use futures::{StreamExt, stream};
use provider::{LlmProvider, build_default_model, require_any_provider_env};
use serde_json::{Value, json};
use std::{
collections::HashMap,
convert::Infallible,
sync::{Arc, Mutex, OnceLock, RwLock},
time::{Duration, Instant},
};
use tracing::{info, warn};
const UI_DEMO_INSTRUCTION: &str = r#"
You are a general-purpose UI assistant.
Design complete decision surfaces, not component inventories. Preserve user-provided facts and use realistic sample data for prototypes.
Prefer high-level tools (`render_layout`, `render_table`, `render_chart`, `render_form`, `render_card`) over raw component assembly.
For dashboards and workspaces, include:
- at least 3 domain-specific KPIs with useful deltas or status
- a populated trend chart or comparison table
- current risks, exceptions, or unresolved decisions
- one clear next action
Represent KPIs as a `render_layout` section with type `stats`; do not encode KPIs as `key_value` pairs.
For intake workflows, provide context, collect only missing information, use specific field labels and options, and end with a review-oriented action instead of a generic Submit button.
Always return complete, responsive surfaces with stable ids. Require review before consequential actions.
"#;
const SUPPORT_INSTRUCTION: &str = r#"
You are a support intake assistant.
On the first turn, render a focused support form that carries known issue facts into its title and description and collects only missing information. A multi-step intake should normally include four to seven purposeful fields chosen from impact summary, affected users or workflow, environment, severity with meaningful options, first observed/deadline, error evidence, and contact preference.
Use concrete placeholders and a domain-specific action such as `Review ticket` or `Review escalation`; never use a generic Submit label.
After form submission, render a review/confirm surface that summarizes the proposed ticket and lets the user cancel. Only after explicit confirmation may you render success with ticket id, SLA, owner, and next step.
Use `render_form` for intake, `render_confirm` for review, and `render_alert` for success/error feedback. Render one surface per turn.
"#;
const APPOINTMENT_INSTRUCTION: &str = r#"
You are a clinic scheduling assistant that renders working UIs.
Use render_layout for overviews and render_page for supporting sections (services, hours, policies).
Use render_card for service options and render_table for schedule availability.
Use render_screen for booking flows and ensure:
- root component id "root"
- layout with Column/Row
- Button actions include action.event.name
After a booking submission, render a confirmation screen with the appointment details.
"#;
const EVENTS_INSTRUCTION: &str = r#"
You are an event RSVP assistant with working UI flows.
Always use render_table for agenda timeline rows with columns: time, session, speaker, room.
Always use render_card for featured speakers with non-empty title and description fields.
Use render_layout for page structure and venue summary sections.
Use render_screen to collect RSVP details (name, guests, dietary, sessions).
Ensure A2UI components include root id "root" and valid Button actions.
Do not compress agenda data into plain text lines.
After submission, render a confirmation screen and a calendar link button.
"#;
const FACILITIES_INSTRUCTION: &str = r#"
You are a facilities maintenance assistant.
Use render_layout for command-center style screens with alerts, KPI cards, and queues.
Use render_table for work-order lists and render_confirm for high-risk actions.
Use render_screen to intake work orders (location, issue type, urgency, contact).
Use render_page for maintenance guidelines or status summaries.
Ensure A2UI components include root id "root".
After intake, render a confirmation with next steps and an emergency contact action.
"#;
const INVENTORY_INSTRUCTION: &str = r#"
You are an inventory restock assistant.
Use render_layout for inventory command views and render_table for stock/reorder grids.
Use render_chart for trend/forecast visuals and render_card for supplier summaries.
Inventory monitor responses must include:
- a stock table with SKU, quantity, threshold, and status
- at least one alert
- at least one chart or summary card row
Use render_screen to collect restock requests (SKU, qty, priority, notes).
Use render_page for inventory summaries and reorder recommendations.
Ensure A2UI components include a root id "root" and explicit child ids.
On submit, show a confirmation card or alert with the request summary.
"#;
const UI_PROTOCOL_HEADER: &str = "x-adk-ui-protocol";
const MAX_REQUEST_BODY_BYTES: usize = 10 * 1024 * 1024;
type RuntimeError = (StatusCode, String);
#[derive(Clone)]
struct ExampleRuntimeState {
config: ServerConfig,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum UiProfile {
AdkUi,
A2ui,
AgUi,
McpApps,
Awp,
}
#[derive(Debug, Clone)]
struct McpAppsNegotiationState {
initialize_request: Option<McpAppsInitializeRequest>,
initialize_result: McpAppsInitializeResult,
initialized: bool,
}
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
struct McpAppsBridgeSessionKey {
app_name: String,
user_id: String,
session_id: String,
}
#[derive(Debug, Clone)]
struct McpAppsBridgeSessionEntry {
initialize_request: Option<McpAppsInitializeRequest>,
initialize_result: McpAppsInitializeResult,
initialized: bool,
message_count: u64,
last_message: Option<Value>,
model_context: Vec<Value>,
model_context_revision: u64,
}
static MCP_APPS_BRIDGE_REGISTRY: OnceLock<
RwLock<HashMap<McpAppsBridgeSessionKey, McpAppsBridgeSessionEntry>>,
> = OnceLock::new();
impl UiProfile {
fn as_str(self) -> &'static str {
match self {
Self::AdkUi => "adk_ui",
Self::A2ui => "a2ui",
Self::AgUi => "ag_ui",
Self::McpApps => "mcp_apps",
Self::Awp => "awp",
}
}
}
fn mcp_apps_bridge_registry()
-> &'static RwLock<HashMap<McpAppsBridgeSessionKey, McpAppsBridgeSessionEntry>> {
MCP_APPS_BRIDGE_REGISTRY.get_or_init(|| RwLock::new(HashMap::new()))
}
fn mcp_apps_bridge_key(app_name: &str, user_id: &str, session_id: &str) -> McpAppsBridgeSessionKey {
McpAppsBridgeSessionKey {
app_name: app_name.to_string(),
user_id: user_id.to_string(),
session_id: session_id.to_string(),
}
}
fn ensure_mcp_apps_bridge_session<'a>(
registry: &'a mut HashMap<McpAppsBridgeSessionKey, McpAppsBridgeSessionEntry>,
app_name: &str,
user_id: &str,
session_id: &str,
) -> &'a mut McpAppsBridgeSessionEntry {
registry
.entry(mcp_apps_bridge_key(app_name, user_id, session_id))
.or_insert_with(|| McpAppsBridgeSessionEntry {
initialize_request: None,
initialize_result: build_default_mcp_apps_initialize_result(
"ui://adk-ui/main",
Some("https://example.com"),
),
initialized: false,
message_count: 0,
last_message: None,
model_context: vec![],
model_context_revision: 0,
})
}
fn negotiation_from_bridge_entry(entry: &McpAppsBridgeSessionEntry) -> McpAppsNegotiationState {
McpAppsNegotiationState {
initialize_request: entry.initialize_request.clone(),
initialize_result: entry.initialize_result.clone(),
initialized: entry.initialized,
}
}
fn full_instruction(extra: &str) -> String {
format!("{A2UI_AGENT_PROMPT}\n\n{extra}")
}
struct SingleSurfaceTurnTool {
inner: Arc<dyn Tool>,
rendered_invocations: Arc<Mutex<HashMap<String, Instant>>>,
}
impl SingleSurfaceTurnTool {
fn new(
inner: Arc<dyn Tool>,
rendered_invocations: Arc<Mutex<HashMap<String, Instant>>>,
) -> Self {
Self {
inner,
rendered_invocations,
}
}
}
#[async_trait]
impl Tool for SingleSurfaceTurnTool {
fn name(&self) -> &str {
self.inner.name()
}
fn description(&self) -> &str {
self.inner.description()
}
fn declaration(&self) -> Value {
self.inner.declaration()
}
async fn execute(&self, ctx: Arc<dyn ToolContext>, args: Value) -> adk_core::Result<Value> {
let invocation_id = ctx.invocation_id().to_string();
let already_rendered = {
let now = Instant::now();
let mut invocations = self
.rendered_invocations
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner());
invocations
.retain(|_, seen_at| now.duration_since(*seen_at) < Duration::from_secs(300));
invocations.insert(invocation_id, now).is_some()
};
if already_rendered {
return Ok(json!({
"status": "ignored",
"reason": "A surface was already rendered for this turn.",
}));
}
let result = self.inner.execute(ctx.clone(), args).await;
if result.is_err() {
let mut invocations = self
.rendered_invocations
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner());
invocations.remove(ctx.invocation_id());
return result;
}
let mut actions = ctx.actions();
actions.skip_summarization = true;
ctx.set_actions(actions);
result
}
}
fn build_ui_agent(
name: &str,
description: &str,
instruction: &str,
model: Arc<dyn Llm>,
ui_tools: &[Arc<dyn Tool>],
) -> Result<Arc<dyn Agent>> {
let mut builder = LlmAgentBuilder::new(name)
.description(description)
.instruction(full_instruction(instruction))
.model(model);
let rendered_invocations = Arc::new(Mutex::new(HashMap::new()));
for tool in ui_tools.iter().cloned() {
builder = builder.tool(Arc::new(SingleSurfaceTurnTool::new(
tool,
rendered_invocations.clone(),
)));
}
Ok(Arc::new(builder.build()?))
}
fn parse_ui_profile(raw: &str) -> Option<UiProfile> {
match normalize_runtime_ui_protocol(raw)? {
"adk_ui" => Some(UiProfile::AdkUi),
"a2ui" => Some(UiProfile::A2ui),
"ag_ui" => Some(UiProfile::AgUi),
"mcp_apps" => Some(UiProfile::McpApps),
"awp" => Some(UiProfile::Awp),
_ => None,
}
}
fn pick_value<'a>(value: &'a Value, keys: &[&str]) -> Option<&'a Value> {
let object = value.as_object()?;
keys.iter().find_map(|key| object.get(*key))
}
fn pick_string(value: &Value, keys: &[&str]) -> Option<String> {
pick_value(value, keys)
.and_then(Value::as_str)
.map(ToOwned::to_owned)
}
fn bridge_body(value: &Value) -> &Value {
pick_value(value, &["params"]).unwrap_or(value)
}
fn bridge_identity(body: &Value) -> Result<(String, String, String), RuntimeError> {
let params = bridge_body(body);
let app_name = pick_string(params, &["appName", "app_name"])
.ok_or_else(|| (StatusCode::BAD_REQUEST, "appName is required".to_string()))?;
let user_id = pick_string(params, &["userId", "user_id"])
.ok_or_else(|| (StatusCode::BAD_REQUEST, "userId is required".to_string()))?;
let session_id = pick_string(params, &["sessionId", "session_id"])
.ok_or_else(|| (StatusCode::BAD_REQUEST, "sessionId is required".to_string()))?;
Ok((app_name, user_id, session_id))
}
fn resolve_ui_profile(headers: &HeaderMap, body: &Value) -> Result<UiProfile, RuntimeError> {
let header_value = headers
.get(UI_PROTOCOL_HEADER)
.and_then(|value| value.to_str().ok());
let body_value = pick_value(body, &["uiProtocol", "ui_protocol"]).and_then(Value::as_str);
let candidate = header_value.or(body_value);
let Some(raw) = candidate else {
return Ok(UiProfile::AdkUi);
};
parse_ui_profile(raw).ok_or_else(|| {
warn!(
requested = %raw,
header = %UI_PROTOCOL_HEADER,
"unsupported ui protocol requested",
);
(
StatusCode::BAD_REQUEST,
format!(
"Unsupported ui protocol '{}'. Supported profiles: {}",
raw,
SUPPORTED_UI_PROTOCOLS.join(", "),
),
)
})
}
fn serialize_runtime_event(event: &Event, profile: UiProfile) -> Option<String> {
if profile == UiProfile::McpApps {
return serialize_runtime_event_with_mcp_bridge(event, profile, None);
}
if profile == UiProfile::AdkUi {
return serde_json::to_string(event).ok();
}
serde_json::to_string(&json!({
"ui_protocol": profile.as_str(),
"event": event,
}))
.ok()
}
fn merge_json(target: &mut Value, overlay: &Value) {
match (target, overlay) {
(Value::Object(target_map), Value::Object(overlay_map)) => {
for (key, value) in overlay_map {
if let Some(existing) = target_map.get_mut(key) {
merge_json(existing, value);
} else {
target_map.insert(key.clone(), value.clone());
}
}
}
(target_slot, overlay_value) => {
*target_slot = overlay_value.clone();
}
}
}
fn looks_like_mcp_apps_payload(value: &Value) -> bool {
let Some(object) = value.as_object() else {
return false;
};
object.contains_key("bridge")
|| object.contains_key("resourceReadResponse")
|| object.contains_key("resource_read_response")
|| object.contains_key("toolResult")
|| object.contains_key("tool_result")
|| object.contains_key("toolMeta")
|| object.contains_key("tool_meta")
}
fn enrich_mcp_apps_payload(value: &mut Value, negotiation: &McpAppsNegotiationState) {
match value {
Value::Array(entries) => {
for entry in entries {
enrich_mcp_apps_payload(entry, negotiation);
}
}
Value::Object(object) => {
for nested in object.values_mut() {
enrich_mcp_apps_payload(nested, negotiation);
}
if !looks_like_mcp_apps_payload(&Value::Object(object.clone())) {
return;
}
let bridge = object
.entry("bridge".to_string())
.or_insert_with(|| Value::Object(Default::default()));
if !bridge.is_object() {
*bridge = Value::Object(Default::default());
}
let bridge_object = bridge.as_object_mut().expect("bridge object");
bridge_object.insert(
"protocolVersion".to_string(),
json!(negotiation.initialize_result.protocol_version),
);
bridge_object.insert(
"hostInfo".to_string(),
serde_json::to_value(&negotiation.initialize_result.host_info)
.unwrap_or_else(|_| json!({})),
);
bridge_object.insert(
"hostCapabilities".to_string(),
negotiation.initialize_result.host_capabilities.clone(),
);
let mut host_context = negotiation.initialize_result.host_context.clone();
if let Some(existing_host_context) = bridge_object.get("hostContext") {
merge_json(&mut host_context, existing_host_context);
}
bridge_object.insert("hostContext".to_string(), host_context);
if let Some(initialize_request) = &negotiation.initialize_request {
bridge_object.insert(
"appInfo".to_string(),
serde_json::to_value(&initialize_request.params.app_info)
.unwrap_or_else(|_| json!({})),
);
let mut app_capabilities = bridge_object
.get("appCapabilities")
.cloned()
.unwrap_or_else(|| json!({}));
let request_capabilities =
serde_json::to_value(&initialize_request.params.app_capabilities)
.unwrap_or_else(|_| json!({}));
merge_json(&mut app_capabilities, &request_capabilities);
bridge_object.insert("appCapabilities".to_string(), app_capabilities);
}
bridge_object.insert("initialized".to_string(), json!(negotiation.initialized));
}
_ => {}
}
}
fn serialize_runtime_event_with_mcp_bridge(
event: &Event,
profile: UiProfile,
negotiation: Option<&McpAppsNegotiationState>,
) -> Option<String> {
let mut serialized = serde_json::to_value(event).ok()?;
if let Some(negotiation) = negotiation {
enrich_mcp_apps_payload(&mut serialized, negotiation);
}
if profile == UiProfile::AdkUi {
return serde_json::to_string(&serialized).ok();
}
serde_json::to_string(&json!({
"ui_protocol": profile.as_str(),
"event": serialized,
}))
.ok()
}
fn log_profile_deprecation(profile: UiProfile) {
if profile != UiProfile::AdkUi {
return;
}
let Some(spec) = UI_PROTOCOL_CAPABILITIES
.iter()
.find(|capability| capability.protocol == profile.as_str())
.and_then(|capability| capability.deprecation)
else {
return;
};
warn!(
protocol = %profile.as_str(),
stage = %spec.stage,
announced_on = %spec.announced_on,
sunset_target_on = ?spec.sunset_target_on,
replacements = ?spec.replacement_protocols,
"legacy ui protocol profile selected",
);
}
fn parse_inline_data(part: &Value) -> Result<Option<Part>, RuntimeError> {
let Some(inline_data) = pick_value(part, &["inlineData", "inline_data"]) else {
return Ok(None);
};
let Some(data) = pick_string(inline_data, &["data"]) else {
return Err((
StatusCode::BAD_REQUEST,
"inlineData.data is required".to_string(),
));
};
let Some(mime_type) = pick_string(inline_data, &["mimeType", "mime_type"]) else {
return Err((
StatusCode::BAD_REQUEST,
"inlineData.mimeType is required".to_string(),
));
};
let decoded = BASE64_STANDARD.decode(data).map_err(|error| {
(
StatusCode::BAD_REQUEST,
format!("Invalid base64 data in inlineData: {}", error),
)
})?;
if decoded.len() > adk_core::MAX_INLINE_DATA_SIZE {
return Err((
StatusCode::PAYLOAD_TOO_LARGE,
format!(
"inlineData exceeds max inline size of {} bytes",
adk_core::MAX_INLINE_DATA_SIZE,
),
));
}
Ok(Some(Part::InlineData {
mime_type,
data: decoded,
uri: None,
annotations: None,
}))
}
fn parse_function_call_part(part: &Value) -> Result<Option<Part>, RuntimeError> {
let Some(function_call) = pick_value(part, &["functionCall", "function_call"]) else {
return Ok(None);
};
let Some(name) = pick_string(function_call, &["name"]) else {
return Err((
StatusCode::BAD_REQUEST,
"functionCall.name is required".to_string(),
));
};
Ok(Some(Part::FunctionCall {
name,
args: pick_value(function_call, &["args", "arguments"])
.cloned()
.unwrap_or(Value::Null),
id: pick_string(function_call, &["id"]),
thought_signature: pick_string(function_call, &["thoughtSignature", "thought_signature"]),
}))
}
fn parse_function_response_part(part: &Value) -> Result<Option<Part>, RuntimeError> {
let Some(function_response) = pick_value(part, &["functionResponse", "function_response"])
else {
return Ok(None);
};
let payload = pick_value(
function_response,
&["functionResponse", "function_response"],
)
.unwrap_or(function_response);
let Some(name) = pick_string(payload, &["name"]) else {
return Err((
StatusCode::BAD_REQUEST,
"functionResponse.name is required".to_string(),
));
};
Ok(Some(Part::FunctionResponse {
function_response: FunctionResponseData {
name,
response: pick_value(payload, &["response", "result", "data", "payload"])
.cloned()
.unwrap_or(Value::Null),
file_data: vec![],
inline_data: vec![],
},
id: pick_string(function_response, &["id"]),
annotations: None,
}))
}
fn parse_message_parts(parts: &[Value]) -> Result<Vec<Part>, RuntimeError> {
let mut parsed = Vec::new();
for part in parts {
if let Some(text) = pick_string(part, &["text"]) {
parsed.push(Part::Text { text });
continue;
}
if let Some(inline_data) = parse_inline_data(part)? {
parsed.push(inline_data);
continue;
}
if let Some(function_call) = parse_function_call_part(part)? {
parsed.push(function_call);
continue;
}
if let Some(function_response) = parse_function_response_part(part)? {
parsed.push(function_response);
}
}
Ok(parsed)
}
fn extract_text_snippets(value: &Value) -> Vec<String> {
match value {
Value::String(text) => {
let trimmed = text.trim();
if trimmed.is_empty() {
Vec::new()
} else {
vec![trimmed.to_string()]
}
}
Value::Array(entries) => entries.iter().flat_map(extract_text_snippets).collect(),
Value::Object(_) => {
if let Some(text) = pick_string(value, &["text"]) {
return vec![text];
}
if let Some(content) = pick_value(value, &["content"]) {
return extract_text_snippets(content);
}
Vec::new()
}
_ => Vec::new(),
}
}
fn ag_ui_input_value(body: &Value) -> Option<&Value> {
pick_value(body, &["input", "agUiInput", "ag_ui_input"])
}
fn ag_ui_activity_id(message: &Value, index: usize) -> String {
pick_string(message, &["id"]).unwrap_or_else(|| {
format!(
"activity-{}-{}",
chrono::Utc::now().timestamp_millis(),
index
)
})
}
fn ag_ui_activity_type(message: &Value) -> String {
pick_string(message, &["activityType", "activity_type", "name"])
.unwrap_or_else(|| "CUSTOM".to_string())
}
fn ag_ui_messages_snapshot_value(body: &Value) -> Option<Value> {
let messages = ag_ui_input_value(body)
.and_then(|input| pick_value(input, &["messages"]))
.and_then(Value::as_array)?;
let filtered: Vec<Value> = messages
.iter()
.filter(|message| {
!(pick_string(message, &["role"]).as_deref() == Some("activity")
&& pick_value(message, &["patch"]).is_some())
})
.cloned()
.collect();
if filtered.is_empty() {
return None;
}
Some(Value::Array(filtered))
}
fn ag_ui_activity_events(body: &Value) -> Vec<Value> {
let Some(messages) = ag_ui_input_value(body)
.and_then(|input| pick_value(input, &["messages"]))
.and_then(Value::as_array)
else {
return Vec::new();
};
messages
.iter()
.enumerate()
.filter(|(_, message)| pick_string(message, &["role"]).as_deref() == Some("activity"))
.map(|(index, message)| {
let message_id = ag_ui_activity_id(message, index);
let activity_type = ag_ui_activity_type(message);
let timestamp = chrono::Utc::now().timestamp_millis().max(0);
if let Some(patch) = pick_value(message, &["patch"]).cloned() {
json!({
"type": "ACTIVITY_DELTA",
"messageId": message_id,
"activityType": activity_type,
"patch": patch,
"timestamp": timestamp,
})
} else {
let content = pick_value(message, &["content"])
.cloned()
.unwrap_or_else(|| json!({}));
let content = match content {
Value::Object(object) => Value::Object(object),
other => json!({ "value": other }),
};
let mut event = json!({
"type": "ACTIVITY_SNAPSHOT",
"messageId": message_id,
"activityType": activity_type,
"content": content,
"timestamp": timestamp,
});
if let Some(replace) = pick_value(message, &["replace"]).and_then(Value::as_bool)
&& let Some(object) = event.as_object_mut()
{
object.insert("replace".to_string(), Value::Bool(replace));
}
event
}
})
.collect()
}
fn a2ui_action_receipt_event(body: &Value) -> Option<String> {
let a2ui_action = pick_value(body, &["a2uiAction", "a2ui_action", "protocolEvent"])?;
let action = pick_value(a2ui_action, &["action"]).unwrap_or(a2ui_action);
let want_response = pick_value(action, &["wantResponse", "want_response"])
.and_then(Value::as_bool)
.unwrap_or(false);
if !want_response {
return None;
}
let action_id = pick_string(action, &["actionId", "action_id"])?;
let name = pick_string(action, &["name"]).unwrap_or_else(|| "action".to_string());
let context = pick_value(action, &["context"])
.cloned()
.unwrap_or(json!({}));
let response = adk_ui::ActionResponseMessage::success(
action_id,
json!({
"received": true,
"status": "accepted_for_processing",
"authoritative": false,
"action": name,
"echoContext": context,
"server": "adk-ui-example",
"protocol": "a2ui",
"version": "v1.0",
}),
);
serde_json::to_string(&response).ok()
}
fn ag_ui_startup_events(body: &Value, session_id: &str) -> Vec<String> {
let Some(ag_ui_input) = ag_ui_input_value(body) else {
return Vec::new();
};
let thread_id = pick_string(ag_ui_input, &["threadId", "thread_id"])
.unwrap_or_else(|| format!("thread-{}", session_id));
let run_id = pick_string(ag_ui_input, &["runId", "run_id"])
.unwrap_or_else(|| format!("run-{}", session_id));
let mut events = vec![json!({
"type": "RUN_STARTED",
"threadId": thread_id,
"runId": run_id,
})];
if let Some(parent_run_id) = pick_string(ag_ui_input, &["parentRunId", "parent_run_id"])
&& let Some(object) = events[0].as_object_mut()
{
object.insert("parentRunId".to_string(), Value::String(parent_run_id));
}
if let Ok(value) = serde_json::to_value(ag_ui_input)
&& let Some(object) = events[0].as_object_mut()
{
object.insert("input".to_string(), value);
}
if let Some(snapshot) = pick_value(ag_ui_input, &["state"]) {
events.push(json!({
"type": "STATE_SNAPSHOT",
"snapshot": snapshot,
}));
}
if let Some(messages) = ag_ui_messages_snapshot_value(body) {
events.push(json!({
"type": "MESSAGES_SNAPSHOT",
"messages": messages,
}));
}
events.extend(ag_ui_activity_events(body));
events
.into_iter()
.filter_map(|event| serde_json::to_string(&event).ok())
.collect()
}
fn build_content_from_request(body: &Value) -> Result<Content, RuntimeError> {
let new_message = pick_value(body, &["newMessage", "new_message"]);
let role = new_message
.and_then(|value| pick_string(value, &["role"]))
.unwrap_or_else(|| "user".to_string());
let mut content = Content::new(role);
if let Some(parts_value) = new_message.and_then(|value| pick_value(value, &["parts"]))
&& let Some(parts) = parts_value.as_array()
{
let parsed_parts = parse_message_parts(parts)?;
if !parsed_parts.is_empty() {
content.parts = parsed_parts;
return Ok(content);
}
}
let mut snippets = Vec::new();
if let Some(ag_ui_input) = pick_value(body, &["input", "agUiInput", "ag_ui_input"])
&& let Some(messages) = pick_value(ag_ui_input, &["messages"]).and_then(Value::as_array)
{
for message in messages {
if pick_string(message, &["role"]).as_deref() == Some("activity") {
continue;
}
if let Some(content_value) = pick_value(message, &["content"]) {
snippets.extend(extract_text_snippets(content_value));
}
}
}
if let Some(mcp_apps_request) = pick_value(body, &["mcpAppsRequest", "mcp_apps_request"])
&& let Some(params) = pick_value(mcp_apps_request, &["params"])
&& let Some(content_value) = pick_value(params, &["content"])
{
snippets.extend(extract_text_snippets(content_value));
}
if let Some(protocol_envelope) = pick_value(body, &["protocolEnvelope", "protocol_envelope"]) {
if let Some(input) = pick_value(protocol_envelope, &["input"])
&& let Some(messages) = pick_value(input, &["messages"]).and_then(Value::as_array)
{
for message in messages {
if let Some(content_value) = pick_value(message, &["content"]) {
snippets.extend(extract_text_snippets(content_value));
}
}
}
if let Some(params) = pick_value(protocol_envelope, &["params"])
&& let Some(content_value) = pick_value(params, &["content"])
{
snippets.extend(extract_text_snippets(content_value));
}
}
for snippet in snippets {
content.parts.push(Part::Text { text: snippet });
}
if content.parts.is_empty() {
return Err((
StatusCode::BAD_REQUEST,
"Request did not contain any renderable message parts.".to_string(),
));
}
Ok(content)
}
fn build_state_delta_from_request(body: &Value, profile: UiProfile) -> HashMap<String, Value> {
let mut state_delta = HashMap::new();
state_delta.insert("ui_protocol".to_string(), json!(profile.as_str()));
if let Some(delta) = pick_value(body, &["stateDelta", "state_delta"]).and_then(Value::as_object)
{
for (key, value) in delta {
state_delta.insert(key.clone(), value.clone());
}
}
if let Some(ui_context) = pick_value(body, &["uiContext", "ui_context"]) {
state_delta.insert("ui_context".to_string(), ui_context.clone());
}
if let Some(ui_event) = pick_value(body, &["uiEvent", "ui_event"]) {
state_delta.insert("ui_event".to_string(), ui_event.clone());
}
if let Some(protocol_event) = pick_value(body, &["protocolEvent", "protocol_event"]) {
state_delta.insert("protocol_event".to_string(), protocol_event.clone());
}
if let Some(protocol_envelope) = pick_value(body, &["protocolEnvelope", "protocol_envelope"]) {
state_delta.insert("protocol_envelope".to_string(), protocol_envelope.clone());
}
if let Some(ag_ui_input) = pick_value(body, &["input", "agUiInput", "ag_ui_input"]) {
if let Some(state) = pick_value(ag_ui_input, &["state"]) {
state_delta.insert("ag_ui_state".to_string(), state.clone());
}
if let Some(forwarded_props) =
pick_value(ag_ui_input, &["forwardedProps", "forwarded_props"])
{
state_delta.insert("ag_ui_forwarded_props".to_string(), forwarded_props.clone());
}
if let Some(thread_id) = pick_value(ag_ui_input, &["threadId", "thread_id"]) {
state_delta.insert("ag_ui_thread_id".to_string(), thread_id.clone());
}
if let Some(run_id) = pick_value(ag_ui_input, &["runId", "run_id"]) {
state_delta.insert("ag_ui_run_id".to_string(), run_id.clone());
}
if let Some(parent_run_id) = pick_value(ag_ui_input, &["parentRunId", "parent_run_id"]) {
state_delta.insert("ag_ui_parent_run_id".to_string(), parent_run_id.clone());
}
if let Some(resume) = pick_value(ag_ui_input, &["resume"]) {
state_delta.insert("ag_ui_resume".to_string(), resume.clone());
}
}
if let Some(compatibility_event) = pick_value(
body,
&["agUiCompatibilityEvent", "ag_ui_compatibility_event"],
) {
state_delta.insert(
"ag_ui_compatibility_event".to_string(),
compatibility_event.clone(),
);
}
if let Some(mcp_apps_request) = pick_value(body, &["mcpAppsRequest", "mcp_apps_request"]) {
state_delta.insert("mcp_apps_request".to_string(), mcp_apps_request.clone());
}
if let Some(mcp_apps_initialize) =
pick_value(body, &["mcpAppsInitialize", "mcp_apps_initialize"])
{
state_delta.insert(
"mcp_apps_initialize".to_string(),
mcp_apps_initialize.clone(),
);
}
if let Some(mcp_apps_initialized) =
pick_value(body, &["mcpAppsInitialized", "mcp_apps_initialized"])
{
state_delta.insert(
"mcp_apps_initialized".to_string(),
mcp_apps_initialized.clone(),
);
}
state_delta
}
fn parse_mcp_apps_initialize_request(
body: &Value,
) -> Result<Option<McpAppsInitializeRequest>, RuntimeError> {
let Some(raw) = pick_value(body, &["mcpAppsInitialize", "mcp_apps_initialize"]) else {
return Ok(None);
};
let request =
serde_json::from_value::<McpAppsInitializeRequest>(raw.clone()).map_err(|error| {
(
StatusCode::BAD_REQUEST,
format!("invalid mcpAppsInitialize payload: {}", error),
)
})?;
if request.method != "ui/initialize" {
return Err((
StatusCode::BAD_REQUEST,
format!(
"unsupported MCP Apps initialize method '{}'",
request.method
),
));
}
Ok(Some(request))
}
fn parse_mcp_apps_initialized(body: &Value) -> Result<bool, RuntimeError> {
let Some(raw) = pick_value(body, &["mcpAppsInitialized", "mcp_apps_initialized"]) else {
return Ok(false);
};
let method = pick_string(raw, &["method"]).ok_or_else(|| {
(
StatusCode::BAD_REQUEST,
"mcpAppsInitialized.method is required".to_string(),
)
})?;
if method != "ui/notifications/initialized" {
return Err((
StatusCode::BAD_REQUEST,
format!("unsupported MCP Apps initialized method '{}'", method),
));
}
Ok(true)
}
fn build_mcp_apps_negotiation_state(
profile: UiProfile,
app_name: &str,
user_id: &str,
session_id: &str,
body: &Value,
) -> Result<Option<McpAppsNegotiationState>, RuntimeError> {
if profile != UiProfile::McpApps {
return Ok(None);
}
let initialized = parse_mcp_apps_initialized(body)?;
let initialize_request = parse_mcp_apps_initialize_request(body)?;
let mut registry = mcp_apps_bridge_registry().write().map_err(|_| {
(
StatusCode::INTERNAL_SERVER_ERROR,
"bridge registry poisoned".to_string(),
)
})?;
let session = ensure_mcp_apps_bridge_session(&mut registry, app_name, user_id, session_id);
if let Some(request) = initialize_request {
session.initialize_result.protocol_version = request.params.protocol_version.clone();
session.initialize_request = Some(request);
}
if initialized {
session.initialized = true;
}
Ok(Some(negotiation_from_bridge_entry(session)))
}
async fn ensure_session(
config: &ServerConfig,
app_name: &str,
user_id: &str,
session_id: &str,
state_delta: &HashMap<String, Value>,
) -> Result<(), RuntimeError> {
let session_exists = config
.session_service
.get(GetRequest {
app_name: app_name.to_string(),
user_id: user_id.to_string(),
session_id: session_id.to_string(),
num_recent_events: None,
after: None,
})
.await
.is_ok();
if !session_exists {
config
.session_service
.create(CreateRequest {
app_name: app_name.to_string(),
user_id: user_id.to_string(),
session_id: Some(session_id.to_string()),
state: state_delta.clone(),
})
.await
.map_err(|error| {
(
StatusCode::INTERNAL_SERVER_ERROR,
format!("failed to create session: {}", error),
)
})?;
return Ok(());
}
if state_delta.is_empty() {
return Ok(());
}
let mut context_event = Event::new(format!("ui-context-{}", session_id));
context_event.author = "user".to_string();
context_event.llm_response.turn_complete = true;
context_event.actions.skip_summarization = true;
context_event.actions.state_delta = state_delta.clone();
config
.session_service
.append_event(session_id, context_event)
.await
.map_err(|error| {
(
StatusCode::INTERNAL_SERVER_ERROR,
format!("failed to persist native ui context: {}", error),
)
})?;
Ok(())
}
fn with_cors_headers(mut response: Response) -> Response {
let headers = response.headers_mut();
headers.insert(
header::ACCESS_CONTROL_ALLOW_ORIGIN,
HeaderValue::from_static("*"),
);
headers.insert(
header::ACCESS_CONTROL_ALLOW_METHODS,
HeaderValue::from_static("GET, POST, PUT, DELETE, OPTIONS"),
);
headers.insert(
header::ACCESS_CONTROL_ALLOW_HEADERS,
HeaderValue::from_static(
"content-type, authorization, x-request-id, x-adk-ui-protocol, x-adk-ui-transport",
),
);
response
}
fn json_error_response(status: StatusCode, message: impl Into<String>) -> Response {
let response = (status, axum::Json(json!({ "error": message.into() }))).into_response();
with_cors_headers(response)
}
fn capabilities_response() -> Response {
#[allow(unused_mut)]
let mut protocols: Vec<serde_json::Value> = UI_PROTOCOL_CAPABILITIES
.iter()
.map(|spec| {
let mut capability = serde_json::to_value(spec).unwrap_or_default();
let details = match spec.protocol {
"a2ui" => json!({
"agentCapabilities": default_a2ui_agent_capabilities(),
"rendererCapabilities": default_a2ui_renderer_capabilities(),
}),
"ag_ui" => json!({
"agentCapabilities": default_ag_ui_agent_capabilities(),
}),
"mcp_apps" => json!({
"protocolVersion": MCP_APPS_PROTOCOL_VERSION,
"extensionCapabilities": default_mcp_apps_extension_capabilities(),
"hostCapabilities": default_mcp_apps_host_capabilities(),
}),
_ => json!({}),
};
if let Some(object) = capability.as_object_mut() {
object.insert("available".to_string(), Value::Bool(true));
object.insert("details".to_string(), details);
}
capability
})
.collect();
#[cfg(feature = "awp")]
if let Ok(mut awp_spec) = serde_json::to_value(&adk_ui::AWP_PROTOCOL_CAPABILITY) {
if let Some(object) = awp_spec.as_object_mut() {
object.insert("available".to_string(), Value::Bool(true));
object.insert("details".to_string(), json!({}));
}
protocols.push(awp_spec);
}
#[allow(unused_mut)]
let mut supported_protocols = SUPPORTED_UI_PROTOCOLS.to_vec();
#[cfg(feature = "awp")]
supported_protocols.push("awp");
let response = axum::Json(json!({
"schemaVersion": UI_CAPABILITY_SCHEMA_VERSION,
"defaultProtocol": UI_DEFAULT_PROTOCOL,
"preferredProtocol": UI_PREFERRED_PROTOCOL,
"supportedProtocols": supported_protocols,
"protocols": protocols,
}))
.into_response();
with_cors_headers(response)
}
fn mcp_apps_initialize_response(body: &Value) -> Result<Response, RuntimeError> {
let params = bridge_body(body);
let (app_name, user_id, session_id) = bridge_identity(body)?;
let protocol_version = pick_string(params, &["protocolVersion", "protocol_version"])
.unwrap_or_else(|| MCP_APPS_PROTOCOL_VERSION.to_string());
let app_info_value = pick_value(params, &["appInfo", "app_info"])
.cloned()
.unwrap_or_else(|| {
json!({
"name": "adk-ui-react-example",
"version": "0.4.0",
})
});
let app_capabilities_value = pick_value(params, &["appCapabilities", "app_capabilities"])
.cloned()
.unwrap_or_else(|| {
json!({
"availableDisplayModes": ["inline"],
"tools": {
"listChanged": false,
}
})
});
let app_info =
serde_json::from_value::<McpAppsPartyInfo>(app_info_value.clone()).map_err(|error| {
(
StatusCode::BAD_REQUEST,
format!("invalid appInfo payload: {}", error),
)
})?;
let app_capabilities = serde_json::from_value::<McpAppsAppCapabilities>(
app_capabilities_value.clone(),
)
.map_err(|error| {
(
StatusCode::BAD_REQUEST,
format!("invalid appCapabilities payload: {}", error),
)
})?;
let mut registry = mcp_apps_bridge_registry().write().map_err(|_| {
(
StatusCode::INTERNAL_SERVER_ERROR,
"bridge registry poisoned".to_string(),
)
})?;
let session = ensure_mcp_apps_bridge_session(&mut registry, &app_name, &user_id, &session_id);
session.initialize_request = Some(McpAppsInitializeRequest {
method: "ui/initialize".to_string(),
params: McpAppsInitializeRequestParams {
protocol_version: protocol_version.clone(),
app_info,
app_capabilities,
},
});
session.initialize_result.protocol_version = protocol_version;
if let Some(host_context) = pick_value(params, &["hostContext", "host_context"]) {
merge_json(&mut session.initialize_result.host_context, host_context);
}
let response = axum::Json(json!({
"initialized": session.initialized,
"protocolVersion": session.initialize_result.protocol_version,
"appInfo": app_info_value,
"appCapabilities": app_capabilities_value,
"hostInfo": serde_json::to_value(&session.initialize_result.host_info).unwrap_or_else(|_| json!({})),
"hostCapabilities": session.initialize_result.host_capabilities,
"hostContext": session.initialize_result.host_context,
"messageCount": session.message_count,
"modelContext": session.model_context,
"modelContextRevision": session.model_context_revision,
}))
.into_response();
Ok(with_cors_headers(response))
}
fn mcp_apps_message_response(body: &Value) -> Result<Response, RuntimeError> {
let params = bridge_body(body);
let (app_name, user_id, session_id) = bridge_identity(body)?;
let mut registry = mcp_apps_bridge_registry().write().map_err(|_| {
(
StatusCode::INTERNAL_SERVER_ERROR,
"bridge registry poisoned".to_string(),
)
})?;
let session = ensure_mcp_apps_bridge_session(&mut registry, &app_name, &user_id, &session_id);
if let Some(host_context) = pick_value(params, &["hostContext", "host_context"]) {
merge_json(&mut session.initialize_result.host_context, host_context);
}
session.message_count += 1;
let mut message = json!({
"role": pick_string(params, &["role"]).unwrap_or_else(|| "user".to_string()),
"content": pick_value(params, &["content"]).cloned().unwrap_or_else(|| json!([])),
});
if let Some(metadata) = pick_value(params, &["metadata"])
&& let Some(object) = message.as_object_mut()
{
object.insert("metadata".to_string(), metadata.clone());
}
session.last_message = Some(message.clone());
let response = axum::Json(json!({
"accepted": true,
"initialized": session.initialized,
"protocolVersion": session.initialize_result.protocol_version,
"messageCount": session.message_count,
"lastMessage": message,
"hostInfo": serde_json::to_value(&session.initialize_result.host_info).unwrap_or_else(|_| json!({})),
"hostCapabilities": session.initialize_result.host_capabilities,
"hostContext": session.initialize_result.host_context,
}))
.into_response();
Ok(with_cors_headers(response))
}
fn mcp_apps_update_model_context_response(body: &Value) -> Result<Response, RuntimeError> {
let params = bridge_body(body);
let (app_name, user_id, session_id) = bridge_identity(body)?;
let mut registry = mcp_apps_bridge_registry().write().map_err(|_| {
(
StatusCode::INTERNAL_SERVER_ERROR,
"bridge registry poisoned".to_string(),
)
})?;
let session = ensure_mcp_apps_bridge_session(&mut registry, &app_name, &user_id, &session_id);
if let Some(host_context) = pick_value(params, &["hostContext", "host_context"]) {
merge_json(&mut session.initialize_result.host_context, host_context);
}
let mode = pick_string(params, &["mode"]).unwrap_or_else(|| "replace".to_string());
let mut next_model_context = if mode.eq_ignore_ascii_case("append") {
session.model_context.clone()
} else {
Vec::new()
};
if let Some(content) = pick_value(params, &["content"]).and_then(Value::as_array) {
next_model_context.extend(content.iter().cloned());
}
if let Some(structured_content) =
pick_value(params, &["structuredContent", "structured_content"])
{
next_model_context.push(json!({
"type": "structuredContent",
"structuredContent": structured_content,
}));
}
session.model_context = next_model_context.clone();
session.model_context_revision += 1;
let response = axum::Json(json!({
"accepted": true,
"initialized": session.initialized,
"protocolVersion": session.initialize_result.protocol_version,
"modelContext": next_model_context,
"modelContextRevision": session.model_context_revision,
"hostInfo": serde_json::to_value(&session.initialize_result.host_info).unwrap_or_else(|_| json!({})),
"hostCapabilities": session.initialize_result.host_capabilities,
"hostContext": session.initialize_result.host_context,
}))
.into_response();
Ok(with_cors_headers(response))
}
async fn handle_protocol_native_run_sse(
state: &ExampleRuntimeState,
headers: &HeaderMap,
body: Value,
) -> Result<Response, RuntimeError> {
let app_name = pick_string(&body, &["appName", "app_name"])
.ok_or_else(|| (StatusCode::BAD_REQUEST, "appName is required".to_string()))?;
let user_id = pick_string(&body, &["userId", "user_id"])
.ok_or_else(|| (StatusCode::BAD_REQUEST, "userId is required".to_string()))?;
let session_id = pick_string(&body, &["sessionId", "session_id"])
.ok_or_else(|| (StatusCode::BAD_REQUEST, "sessionId is required".to_string()))?;
let streaming = pick_value(&body, &["streaming"])
.and_then(Value::as_bool)
.unwrap_or(true);
let ui_profile = resolve_ui_profile(headers, &body)?;
let mcp_apps_negotiation =
build_mcp_apps_negotiation_state(ui_profile, &app_name, &user_id, &session_id, &body)?;
log_profile_deprecation(ui_profile);
info!(
app_name = %app_name,
user_id = %user_id,
session_id = %session_id,
ui_protocol = %ui_profile.as_str(),
"protocol-aware /api/run_sse request received",
);
let content = build_content_from_request(&body)?;
let state_delta = build_state_delta_from_request(&body, ui_profile);
ensure_session(
&state.config,
&app_name,
&user_id,
&session_id,
&state_delta,
)
.await?;
let agent = state
.config
.agent_loader
.load_agent(&app_name)
.await
.map_err(|error| {
(
StatusCode::INTERNAL_SERVER_ERROR,
format!("failed to load agent: {}", error),
)
})?;
let streaming_mode = if streaming {
adk_core::StreamingMode::SSE
} else {
adk_core::StreamingMode::None
};
let runner = Runner::new(RunnerConfig {
app_name: app_name.clone(),
agent,
session_service: state.config.session_service.clone(),
artifact_service: state.config.artifact_service.clone(),
memory_service: state.config.memory_service.clone(),
plugin_manager: None,
run_config: Some(adk_core::RunConfig {
streaming_mode,
..adk_core::RunConfig::default()
}),
compaction_config: state.config.compaction_config.clone(),
context_cache_config: state.config.context_cache_config.clone(),
cache_capable: state.config.cache_capable.clone(),
request_context: None,
cancellation_token: None,
intra_compaction_config: None,
intra_compaction_summarizer: None,
})
.map_err(|error| {
(
StatusCode::INTERNAL_SERVER_ERROR,
format!("failed to create runner: {}", error),
)
})?;
let event_stream = runner
.run(
UserId::try_from(user_id.as_str())
.map_err(|e| (StatusCode::BAD_REQUEST, format!("invalid userId: {}", e)))?,
SessionId::try_from(session_id.as_str())
.map_err(|e| (StatusCode::BAD_REQUEST, format!("invalid sessionId: {}", e)))?,
content,
)
.await
.map_err(|error| {
(
StatusCode::INTERNAL_SERVER_ERROR,
format!("failed to run agent: {}", error),
)
})?;
let selected_profile = ui_profile;
let mcp_bridge_state = mcp_apps_negotiation.clone();
let mut startup_events: Vec<String> = Vec::new();
if selected_profile == UiProfile::AgUi {
startup_events.extend(ag_ui_startup_events(&body, &session_id));
}
if let Some(response_json) = a2ui_action_receipt_event(&body) {
startup_events.push(response_json);
}
let startup_stream = stream::iter(
startup_events
.into_iter()
.map(|json| Ok::<_, Infallible>(SseEvent::default().data(json))),
);
let runtime_stream = event_stream.filter_map(move |item| {
let mcp_bridge_state = mcp_bridge_state.clone();
async move {
match item {
Ok(event) => {
let json = if selected_profile == UiProfile::McpApps {
serialize_runtime_event_with_mcp_bridge(
&event,
selected_profile,
mcp_bridge_state.as_ref(),
)?
} else {
serialize_runtime_event(&event, selected_profile)?
};
Some(Ok::<_, Infallible>(SseEvent::default().data(json)))
}
Err(_) => None,
}
}
});
let sse_stream = startup_stream.chain(runtime_stream);
let response = Sse::new(sse_stream)
.keep_alive(KeepAlive::default())
.into_response();
Ok(with_cors_headers(response))
}
async fn protocol_native_run_sse_middleware(
State(state): State<ExampleRuntimeState>,
request: Request<Body>,
next: Next,
) -> Response {
let is_options = request.method() == Method::OPTIONS;
let path = request.uri().path().to_string();
let should_intercept_run_sse = request.method() == Method::POST && path == "/api/run_sse";
let should_intercept_capabilities =
request.method() == Method::GET && path == "/api/ui/capabilities";
let should_intercept_mcp_initialize =
request.method() == Method::POST && path == "/api/ui/initialize";
let should_intercept_mcp_message =
request.method() == Method::POST && path == "/api/ui/message";
let should_intercept_mcp_model_context =
request.method() == Method::POST && path == "/api/ui/update-model-context";
let should_intercept_preflight = is_options
&& matches!(
path.as_str(),
"/api/run_sse"
| "/api/ui/capabilities"
| "/api/ui/initialize"
| "/api/ui/message"
| "/api/ui/update-model-context"
);
if should_intercept_preflight {
return with_cors_headers(StatusCode::OK.into_response());
}
if should_intercept_capabilities {
return capabilities_response();
}
if !should_intercept_run_sse
&& !should_intercept_mcp_initialize
&& !should_intercept_mcp_message
&& !should_intercept_mcp_model_context
{
return next.run(request).await;
}
let (parts, body) = request.into_parts();
let headers = parts.headers.clone();
let body_bytes = match to_bytes(body, MAX_REQUEST_BODY_BYTES).await {
Ok(bytes) => bytes,
Err(error) => {
return json_error_response(
StatusCode::BAD_REQUEST,
format!("failed to read request body: {}", error),
);
}
};
let body_value = match serde_json::from_slice::<Value>(&body_bytes) {
Ok(value) => value,
Err(error) => {
return json_error_response(
StatusCode::BAD_REQUEST,
format!("invalid request body: {}", error),
);
}
};
if should_intercept_mcp_initialize {
return match mcp_apps_initialize_response(&body_value) {
Ok(response) => response,
Err((status, message)) => json_error_response(status, message),
};
}
if should_intercept_mcp_message {
return match mcp_apps_message_response(&body_value) {
Ok(response) => response,
Err((status, message)) => json_error_response(status, message),
};
}
if should_intercept_mcp_model_context {
return match mcp_apps_update_model_context_response(&body_value) {
Ok(response) => response,
Err((status, message)) => json_error_response(status, message),
};
}
match handle_protocol_native_run_sse(&state, &headers, body_value).await {
Ok(response) => response,
Err((status, message)) => json_error_response(status, message),
}
}
const AGENTIC_A2UI_INSTRUCTION: &str = r#"
You are an agentic support intake assistant that emits A2UI.
Always set protocol=a2ui on render tools when the option exists.
Render exactly one user-facing surface per turn. Never call several render tools in the same turn.
Choose the surface from the current interaction state:
1) If required issue details are missing, render_form and stop. Carry known facts into the title and description, then collect four to seven purposeful missing details such as impact, affected users/workflow, environment, severity, deadline, evidence, and contact preference. Use meaningful select options and a review-oriented CTA, never a generic Submit label.
2) If the issue details are present but escalation is not confirmed, render_confirm and stop.
3) Only after an explicit confirm UI event, render_toast or render_alert and stop.
Never render a success state before the user confirms. A high-level render tool already returns a complete surface, so do not follow it with a placeholder render_screen call.
Preserve submitted state in the review and make the escalation scope clear before confirmation.
Prefer render_screen/render_page with root id "root" and stable component ids.
Never dump raw JSON for the user — always render UI surfaces.
"#;
const AGENTIC_AG_UI_INSTRUCTION: &str = r#"
You are an operations agent for AG-UI clients.
Use protocol=ag_ui on render tools. Treat the UI as a live decision workspace and chat text as secondary.
When the request describes a complete operational product with distinct routes, use render_app with two to five purposeful pages instead of compressing every job into one workspace.
Render populated operational context: impact KPIs, service or queue health, recent changes, evidence timeline, owners, current investigation state, and the recommended next action. Use realistic sample values when live data is unavailable and distinguish them from user-provided facts.
Use render_layout for the complete workspace. When an operator must edit a rollback or remediation proposal, include it as a `form` section in that layout; never compress the workspace into one long render_form description. Keep the workspace to roughly six to eight sections by combining related records.
Preserve shared state across updates. Surface progress and concise reasoning status without exposing private chain-of-thought.
For failover, rollback, or other risky actions, present an editable proposal and require explicit approval before rendering completion.
"#;
const AGENTIC_MCP_APPS_INSTRUCTION: &str = r#"
You are a host-embedded MCP Apps agent.
Use protocol=mcp_apps on render tools so results include ui:// resources and structured content.
Create a focused interactive app that preserves host context and user state across turns. Use a compact card for one decision and a multi-section layout for filtering, inspection, comparison, or editing.
When the request describes a complete host-embedded product or multi-route workflow, use render_app with two to five purposeful pages.
Keep resource URIs stable (ui://adk-ui/...), use domain-specific populated content, and do not reproduce controls already owned by the host conversation.
Limit a focused card to one primary and one secondary action. Require an editable review/confirm step before privileged or destructive actions.
"#;
#[tokio::main]
async fn main() -> Result<()> {
dotenvy::dotenv().ok();
require_any_provider_env()?;
let (model, model_name, provider) = build_default_model()?;
let ui_tools = UiToolset::all_tools();
let ui_demo = build_ui_agent(
"ui_demo",
"General purpose multi-surface demo agent",
UI_DEMO_INSTRUCTION,
model.clone(),
&ui_tools,
)?;
let ui_working_support = build_ui_agent(
"ui_working_support",
"Support intake agent with working UI flows",
SUPPORT_INSTRUCTION,
model.clone(),
&ui_tools,
)?;
let ui_working_appointment = build_ui_agent(
"ui_working_appointment",
"Appointment scheduling agent with working UI flows",
APPOINTMENT_INSTRUCTION,
model.clone(),
&ui_tools,
)?;
let ui_working_events = build_ui_agent(
"ui_working_events",
"Event RSVP agent with working UI flows",
EVENTS_INSTRUCTION,
model.clone(),
&ui_tools,
)?;
let ui_working_facilities = build_ui_agent(
"ui_working_facilities",
"Facilities maintenance agent with working UI flows",
FACILITIES_INSTRUCTION,
model.clone(),
&ui_tools,
)?;
let ui_working_inventory = build_ui_agent(
"ui_working_inventory",
"Inventory restock agent with working UI flows",
INVENTORY_INSTRUCTION,
model.clone(),
&ui_tools,
)?;
let agentic_a2ui = build_ui_agent(
"agentic_a2ui_support",
"Agentic A2UI support intake multi-step flow",
AGENTIC_A2UI_INSTRUCTION,
model.clone(),
&ui_tools,
)?;
let agentic_ag_ui = build_ui_agent(
"agentic_ag_ui_ops",
"Agentic AG-UI operations command center",
AGENTIC_AG_UI_INSTRUCTION,
model.clone(),
&ui_tools,
)?;
let agentic_mcp_apps = build_ui_agent(
"agentic_mcp_apps_confirm",
"Agentic MCP Apps confirm / form host bridge demo",
AGENTIC_MCP_APPS_INSTRUCTION,
model,
&ui_tools,
)?;
let agent_loader = Arc::new(MultiAgentLoader::new(vec![
ui_demo,
ui_working_support,
ui_working_appointment,
ui_working_events,
ui_working_facilities,
ui_working_inventory,
agentic_a2ui,
agentic_ag_ui,
agentic_mcp_apps,
])?);
let port = std::env::var("PORT")
.ok()
.and_then(|value| value.parse().ok())
.unwrap_or(8080);
let session_service: Arc<dyn SessionService> = Arc::new(InMemorySessionService::new());
let config = ServerConfig::new(agent_loader, session_service);
let app = create_app(config.clone()).layer(middleware::from_fn_with_state(
ExampleRuntimeState { config },
protocol_native_run_sse_middleware,
));
let listener = tokio::net::TcpListener::bind(format!("0.0.0.0:{port}")).await?;
println!("=== ADK UI Aggregated Server ===");
println!("Server running on http://localhost:{}", port);
println!(
"Provider: {} (set ADK_UI_PROVIDER=openai|gemini)",
provider.as_str()
);
println!("Model: {}", model_name);
if provider == LlmProvider::OpenAi {
println!("Default is OpenAI — use OPENAI_API_KEY / ADK_UI_MODEL");
} else {
println!("Gemini selected — use GOOGLE_API_KEY / ADK_UI_MODEL");
}
println!();
println!("Loaded apps:");
println!(" - ui_demo");
println!(" - ui_working_support");
println!(" - ui_working_appointment");
println!(" - ui_working_events");
println!(" - ui_working_facilities");
println!(" - ui_working_inventory");
println!(" - agentic_a2ui_support (A2UI multi-step intake)");
println!(" - agentic_ag_ui_ops (AG-UI ops command center)");
println!(" - agentic_mcp_apps_confirm (MCP Apps host bridge)");
println!();
println!("React client: http://localhost:5173");
println!();
axum::serve(listener, app)
.with_graceful_shutdown(shutdown_signal())
.await?;
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn build_content_falls_back_to_ag_ui_messages() {
let body = json!({
"agUiInput": {
"messages": [
{ "content": "Approve the request" }
]
}
});
let content = build_content_from_request(&body).expect("ag-ui content");
assert_eq!(content.role, "user");
assert_eq!(content.parts.len(), 1);
assert!(matches!(&content.parts[0], Part::Text { text } if text == "Approve the request"));
}
#[test]
fn build_content_falls_back_to_ag_ui_input_alias() {
let body = json!({
"input": {
"messages": [
{
"content": [
{ "type": "text", "text": "Review the dashboard" }
]
}
]
}
});
let content = build_content_from_request(&body).expect("ag-ui alias content");
assert_eq!(content.role, "user");
assert_eq!(content.parts.len(), 1);
assert!(matches!(&content.parts[0], Part::Text { text } if text == "Review the dashboard"));
}
#[test]
fn build_content_supports_function_response_parts() {
let body = json!({
"newMessage": {
"role": "user",
"parts": [
{
"functionResponse": {
"name": "approve_request",
"response": { "approved": true }
}
}
]
}
});
let content = build_content_from_request(&body).expect("function response content");
assert_eq!(content.parts.len(), 1);
assert!(matches!(
&content.parts[0],
Part::FunctionResponse {
function_response: FunctionResponseData { name, response },
..
} if name == "approve_request" && response == &json!({ "approved": true })
));
}
#[test]
fn build_state_delta_captures_native_protocol_context() {
let body = json!({
"uiContext": {
"surfaceId": "main"
},
"uiEvent": {
"action": "button_click",
"action_id": "approve_request"
},
"agUiInput": {
"state": {
"adkUi": {
"surfaceId": "main"
}
},
"forwardedProps": {
"source": "adk-ui-react"
},
"threadId": "thread-123",
"runId": "run-456",
"resume": [{
"interruptId": "interrupt-1",
"status": "resolved",
"payload": {"approved": true}
}]
}
});
let state_delta = build_state_delta_from_request(&body, UiProfile::AgUi);
assert_eq!(state_delta.get("ui_protocol"), Some(&json!("ag_ui")));
assert_eq!(
state_delta.get("ui_context"),
Some(&json!({ "surfaceId": "main" })),
);
assert_eq!(
state_delta.get("ag_ui_thread_id"),
Some(&json!("thread-123")),
);
assert_eq!(state_delta.get("ag_ui_run_id"), Some(&json!("run-456")),);
assert_eq!(
state_delta.get("ag_ui_resume"),
Some(&json!([{
"interruptId": "interrupt-1",
"status": "resolved",
"payload": {"approved": true}
}])),
);
}
#[test]
fn build_mcp_apps_negotiation_state_parses_initialize_request() {
let body = json!({
"mcpAppsInitialize": {
"method": "ui/initialize",
"params": {
"protocolVersion": "2025-11-21",
"appInfo": {
"name": "adk-ui-react-example",
"version": "0.4.0",
},
"appCapabilities": {
"availableDisplayModes": ["inline"],
"tools": {
"listChanged": false,
}
}
}
},
"mcpAppsInitialized": {
"method": "ui/notifications/initialized",
"params": {}
}
});
let negotiation = build_mcp_apps_negotiation_state(
UiProfile::McpApps,
"ui_demo",
"user1",
"session-1",
&body,
)
.expect("negotiation");
let negotiation = negotiation.expect("mcp apps negotiation state");
assert!(negotiation.initialized);
assert_eq!(
negotiation
.initialize_request
.as_ref()
.expect("initialize request")
.params
.app_info
.name,
"adk-ui-react-example"
);
assert_eq!(negotiation.initialize_result.protocol_version, "2025-11-21");
}
#[test]
fn mcp_apps_direct_initialize_is_reused_by_runtime_negotiation() {
let body = json!({
"appName": "ui_demo",
"userId": "user1",
"sessionId": "session-direct",
"protocolVersion": "2025-11-21",
"appInfo": {
"name": "adk-ui-react-example",
"version": "0.4.0"
},
"appCapabilities": {
"availableDisplayModes": ["inline"]
}
});
let _ = mcp_apps_initialize_response(&body).expect("initialize response");
let negotiation = build_mcp_apps_negotiation_state(
UiProfile::McpApps,
"ui_demo",
"user1",
"session-direct",
&json!({
"mcpAppsInitialized": {
"method": "ui/notifications/initialized"
}
}),
)
.expect("negotiation")
.expect("mcp apps negotiation state");
assert!(negotiation.initialized);
assert_eq!(
negotiation
.initialize_request
.as_ref()
.expect("initialize request")
.params
.app_info
.name,
"adk-ui-react-example"
);
}
}