From 0ee99c6d9c60dcd84c03fee377ee26fdb82ac89b Mon Sep 17 00:00:00 2001
From: David Soria Parra <167242713+dsp-ant@users.noreply.github.com>
Date: Wed, 4 Dec 2024 15:45:25 +0000
Subject: [PATCH] context_server: Add missing types for MCP spec to protocol
 2024-11-05 (#21498)

This commit syncs missing types for the mcp spec 2024-11-05.

Release Notes:

- N/A
---
 .../slash_command/context_server_command.rs   |  2 +-
 crates/context_server/src/types.rs            | 82 ++++++++++++++++++-
 2 files changed, 79 insertions(+), 5 deletions(-)

diff --git a/crates/assistant/src/slash_command/context_server_command.rs b/crates/assistant/src/slash_command/context_server_command.rs
index b183a77f544e7..8c53ddb773eb5 100644
--- a/crates/assistant/src/slash_command/context_server_command.rs
+++ b/crates/assistant/src/slash_command/context_server_command.rs
@@ -164,7 +164,7 @@ impl SlashCommand for ContextServerSlashCommand {
                     .messages
                     .into_iter()
                     .filter_map(|msg| match msg.content {
-                        context_server::types::MessageContent::Text { text } => Some(text),
+                        context_server::types::MessageContent::Text { text, .. } => Some(text),
                         _ => None,
                     })
                     .collect::<Vec<String>>()
diff --git a/crates/context_server/src/types.rs b/crates/context_server/src/types.rs
index 851ebbf08ba94..f3c6e1c5e2e1b 100644
--- a/crates/context_server/src/types.rs
+++ b/crates/context_server/src/types.rs
@@ -167,11 +167,18 @@ pub struct InitializeResponse {
 #[derive(Debug, Deserialize)]
 #[serde(rename_all = "camelCase")]
 pub struct ResourcesReadResponse {
-    pub contents: Vec<ResourceContents>,
+    pub contents: Vec<ResourceContentsType>,
     #[serde(rename = "_meta", skip_serializing_if = "Option::is_none")]
     pub meta: Option<HashMap<String, serde_json::Value>>,
 }
 
+#[derive(Debug, Deserialize)]
+#[serde(untagged)]
+pub enum ResourceContentsType {
+    Text(TextResourceContents),
+    Blob(BlobResourceContents),
+}
+
 #[derive(Debug, Deserialize)]
 #[serde(rename_all = "camelCase")]
 pub struct ResourcesListResponse {
@@ -181,6 +188,7 @@ pub struct ResourcesListResponse {
     #[serde(rename = "_meta", skip_serializing_if = "Option::is_none")]
     pub meta: Option<HashMap<String, serde_json::Value>>,
 }
+
 #[derive(Debug, Serialize, Deserialize)]
 #[serde(rename_all = "camelCase")]
 pub struct SamplingMessage {
@@ -188,6 +196,35 @@ pub struct SamplingMessage {
     pub content: MessageContent,
 }
 
+#[derive(Debug, Serialize)]
+#[serde(rename_all = "camelCase")]
+pub struct CreateMessageRequest {
+    pub messages: Vec<SamplingMessage>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub model_preferences: Option<ModelPreferences>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub system_prompt: Option<String>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub include_context: Option<String>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub temperature: Option<f64>,
+    pub max_tokens: u32,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub stop_sequences: Option<Vec<String>>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub metadata: Option<serde_json::Value>,
+}
+
+#[derive(Debug, Deserialize)]
+#[serde(rename_all = "camelCase")]
+pub struct CreateMessageResult {
+    pub role: Role,
+    pub content: MessageContent,
+    pub model: String,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub stop_reason: Option<String>,
+}
+
 #[derive(Debug, Serialize, Deserialize)]
 #[serde(rename_all = "camelCase")]
 pub struct PromptMessage {
@@ -206,11 +243,33 @@ pub enum Role {
 #[serde(tag = "type")]
 pub enum MessageContent {
     #[serde(rename = "text")]
-    Text { text: String },
+    Text {
+        text: String,
+        #[serde(skip_serializing_if = "Option::is_none")]
+        annotations: Option<MessageAnnotations>,
+    },
     #[serde(rename = "image")]
-    Image { data: String, mime_type: String },
+    Image {
+        data: String,
+        mime_type: String,
+        #[serde(skip_serializing_if = "Option::is_none")]
+        annotations: Option<MessageAnnotations>,
+    },
     #[serde(rename = "resource")]
-    Resource { resource: ResourceContents },
+    Resource {
+        resource: ResourceContents,
+        #[serde(skip_serializing_if = "Option::is_none")]
+        annotations: Option<MessageAnnotations>,
+    },
+}
+
+#[derive(Debug, Serialize, Deserialize)]
+#[serde(rename_all = "camelCase")]
+pub struct MessageAnnotations {
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub audience: Option<Vec<Role>>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub priority: Option<f64>,
 }
 
 #[derive(Debug, Deserialize)]
@@ -460,6 +519,11 @@ pub enum ClientNotification {
     Initialized,
     Progress(ProgressParams),
     RootsListChanged,
+    Cancelled {
+        request_id: String,
+        #[serde(skip_serializing_if = "Option::is_none")]
+        reason: Option<String>,
+    },
 }
 
 #[derive(Debug, Serialize, Deserialize)]
@@ -532,6 +596,16 @@ pub struct ListToolsResponse {
     pub meta: Option<HashMap<String, serde_json::Value>>,
 }
 
+#[derive(Debug, Deserialize)]
+#[serde(rename_all = "camelCase")]
+pub struct ListResourceTemplatesResponse {
+    pub resource_templates: Vec<ResourceTemplate>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub next_cursor: Option<String>,
+    #[serde(rename = "_meta", skip_serializing_if = "Option::is_none")]
+    pub meta: Option<HashMap<String, serde_json::Value>>,
+}
+
 #[derive(Debug, Deserialize)]
 #[serde(rename_all = "camelCase")]
 pub struct ListRootsResponse {