#![cfg(not(target_arch = "wasm32"))]
use std::net::SocketAddr;
use std::sync::Arc;
use pmcp::server::streamable_http_server::StreamableHttpServer;
use pmcp::server::task_store::{InMemoryTaskStore, TaskStore};
use pmcp::server::typed_tool::TypedTool;
use pmcp::shared::streamable_http::{StreamableHttpTransport, StreamableHttpTransportConfig};
use pmcp::types::tasks::TaskMetadata;
use pmcp::types::{CallToolResult, Content};
use pmcp::{Client, Error, Server};
use schemars::JsonSchema;
use serde::Deserialize;
use tokio::sync::Mutex;
use url::Url;
#[derive(Debug, Deserialize, JsonSchema)]
struct StartArgs {}
#[tokio::main]
async fn main() -> pmcp::Result<()> {
let store = Arc::new(InMemoryTaskStore::new());
let minted = store
.create("s47-owner", None)
.await
.map_err(|e| Error::internal(format!("store.create failed: {e}")))?;
let after_task_id = minted.task_id.clone();
println!("store-minted task id (AFTER points here): {after_task_id}");
let closure_task_id = after_task_id.clone();
let before_tool = TypedTool::new_with_schema(
"before_hand_rolled",
serde_json::json!({ "type": "object" }),
|_args: serde_json::Value, _extra| {
Box::pin(async {
Ok(serde_json::json!({
"content": [
{ "type": "text", "text": "processing started (hand-rolled)" }
],
"_meta": {
"io.modelcontextprotocol/related-task": {
"taskId": "hand-rolled-not-store-minted"
}
}
}))
})
},
)
.with_description("BEFORE: hand-rolls _meta on the Payload path (anti-pattern)");
let server = Server::builder()
.name("s47-task-augmented-result")
.version("1.0.0")
.tool("before_hand_rolled", before_tool)
.suppress_double_wrap_check("before_hand_rolled")
.tool_with_result::<StartArgs>("after_native", move |_args, _extra| {
let task_id = closure_task_id.clone();
Box::pin(async move {
Ok(
CallToolResult::new(vec![Content::text("processing started (native)")])
.with_related_task(TaskMetadata::new(task_id)),
)
})
})
.task_store(store.clone() as Arc<dyn TaskStore>)
.build()?;
let server = Arc::new(Mutex::new(server));
let bind_addr: SocketAddr = "127.0.0.1:0".parse().expect("valid loopback addr");
let http_server = StreamableHttpServer::new(bind_addr, server);
let (bound, server_handle) = http_server.start().await?;
println!("HTTP server listening on http://{bound}");
let result = run_migration_demo(bound, &after_task_id).await;
server_handle.abort();
println!("server task aborted (clean shutdown)");
result
}
async fn run_migration_demo(bound: SocketAddr, expected_task_id: &str) -> pmcp::Result<()> {
let config = StreamableHttpTransportConfig {
url: Url::parse(&format!("http://{bound}")).map_err(|e| Error::Internal(e.to_string()))?,
extra_headers: vec![],
auth_provider: None,
session_id: None,
enable_json_response: true,
on_resumption_token: None,
http_middleware_chain: None,
};
let mut client = Client::new(StreamableHttpTransport::new(config));
client
.initialize(pmcp::types::ClientCapabilities::default())
.await?;
let before = client
.call_tool("before_hand_rolled".to_string(), serde_json::json!({}))
.await?;
println!("\nBEFORE (hand-rolled):");
println!(" related_task() = {:?}", before.related_task());
if before.related_task().is_some() {
return Err(Error::internal(
"BEFORE demo invariant broken: the hand-rolled _meta unexpectedly survived at \
result top level — the anti-pattern would not be demonstrated",
));
}
let before_text = before
.content
.first()
.and_then(|c| match c {
Content::Text { text } => Some(text.clone()),
_ => None,
})
.unwrap_or_default();
if !before_text.contains("_meta") {
return Err(Error::internal(
"BEFORE demo invariant broken: expected the hand-rolled envelope (incl. its _meta) \
to be stringified into content[0].text",
));
}
println!(" content[0].text contains the stringified _meta envelope (the double-wrap bug)");
let after = client
.call_tool("after_native".to_string(), serde_json::json!({}))
.await?;
println!("\nAFTER (native task_store + tool_with_result):");
let related = after.related_task().ok_or_else(|| {
Error::internal("AFTER result must expose related_task() == Some over the wire")
})?;
println!(" related_task().taskId = {}", related.task_id);
if related.task_id != expected_task_id {
return Err(Error::internal(format!(
"AFTER related-task id must equal the store-minted id ({expected_task_id}), got {}",
related.task_id
)));
}
let after_text = after
.content
.first()
.and_then(|c| match c {
Content::Text { text } => Some(text.clone()),
_ => None,
})
.unwrap_or_default();
if after_text.contains("_meta") {
return Err(Error::internal(
"AFTER content must NOT be a stringified envelope (no _meta in content[0].text)",
));
}
println!(" content[0].text = {after_text:?} (not stringified)");
println!(
"\nMigration OK: BEFORE loses _meta (anti-pattern), AFTER carries the store-minted \
related-task at result top level (native pattern)."
);
Ok(())
}