Update to new agent schema (#35578)

Release Notes:

- N/A

---------

Co-authored-by: Agus Zubiaga <agus@zed.dev>

crates/context_server/src/client.rs

@@ -441,14 +441,12 @@ impl Client {
         Ok(())
     }
 
-    #[allow(unused)]
-    pub fn on_notification<F>(&self, method: &'static str, f: F)
-    where
-        F: 'static + Send + FnMut(Value, AsyncApp),
-    {
-        self.notification_handlers
-            .lock()
-            .insert(method, Box::new(f));
+    pub fn on_notification(
+        &self,
+        method: &'static str,
+        f: Box<dyn 'static + Send + FnMut(Value, AsyncApp)>,
+    ) {
+        self.notification_handlers.lock().insert(method, f);
     }
 }
 

crates/context_server/src/context_server.rs

@@ -95,8 +95,28 @@ impl ContextServer {
         self.client.read().clone()
     }
 
-    pub async fn start(self: Arc<Self>, cx: &AsyncApp) -> Result<()> {
-        let client = match &self.configuration {
+    pub async fn start(&self, cx: &AsyncApp) -> Result<()> {
+        self.initialize(self.new_client(cx)?).await
+    }
+
+    /// Starts the context server, making sure handlers are registered before initialization happens
+    pub async fn start_with_handlers(
+        &self,
+        notification_handlers: Vec<(
+            &'static str,
+            Box<dyn 'static + Send + FnMut(serde_json::Value, AsyncApp)>,
+        )>,
+        cx: &AsyncApp,
+    ) -> Result<()> {
+        let client = self.new_client(cx)?;
+        for (method, handler) in notification_handlers {
+            client.on_notification(method, handler);
+        }
+        self.initialize(client).await
+    }
+
+    fn new_client(&self, cx: &AsyncApp) -> Result<Client> {
+        Ok(match &self.configuration {
             ContextServerTransport::Stdio(command, working_directory) => Client::stdio(
                 client::ContextServerId(self.id.0.clone()),
                 client::ModelContextServerBinary {
@@ -113,8 +133,7 @@ impl ContextServer {
                 transport.clone(),
                 cx.clone(),
             )?,
-        };
-        self.initialize(client).await
+        })
     }
 
     async fn initialize(&self, client: Client) -> Result<()> {

crates/context_server/src/listener.rs

@@ -83,14 +83,18 @@ impl McpServer {
     }
 
     pub fn add_tool<T: McpServerTool + Clone + 'static>(&mut self, tool: T) {
-        let output_schema = schemars::schema_for!(T::Output);
-        let unit_schema = schemars::schema_for!(());
+        let mut settings = schemars::generate::SchemaSettings::draft07();
+        settings.inline_subschemas = true;
+        let mut generator = settings.into_generator();
+
+        let output_schema = generator.root_schema_for::<T::Output>();
+        let unit_schema = generator.root_schema_for::<T::Output>();
 
         let registered_tool = RegisteredTool {
             tool: Tool {
                 name: T::NAME.into(),
                 description: Some(tool.description().into()),
-                input_schema: schemars::schema_for!(T::Input).into(),
+                input_schema: generator.root_schema_for::<T::Input>().into(),
                 output_schema: if output_schema == unit_schema {
                     None
                 } else {

crates/context_server/src/protocol.rs

@@ -115,10 +115,11 @@ impl InitializedContextServerProtocol {
         self.inner.notify(T::METHOD, params)
     }
 
-    pub fn on_notification<F>(&self, method: &'static str, f: F)
-    where
-        F: 'static + Send + FnMut(Value, AsyncApp),
-    {
+    pub fn on_notification(
+        &self,
+        method: &'static str,
+        f: Box<dyn 'static + Send + FnMut(Value, AsyncApp)>,
+    ) {
         self.inner.on_notification(method, f);
     }
 }