Lifei/acp session setup refactor (#9488)

crates/goose-sdk/src/custom_notifications.rs

@@ -0,0 +1,173 @@
+use crate::custom_requests::CustomMethodSchema;
+use agent_client_protocol::{JsonRpcMessage, JsonRpcNotification};
+use schemars::{JsonSchema, SchemaGenerator};
+use serde::{Deserialize, Serialize};
+
+/// Goose-custom session update notification — a parallel to ACP's
+/// `session/update` carrying goose-specific update variants.
+#[derive(Debug, Default, Clone, Serialize, Deserialize, JsonSchema, JsonRpcNotification)]
+#[notification(method = "_goose/unstable/session/update")]
+#[serde(rename_all = "camelCase")]
+pub struct GooseSessionNotification {
+    pub session_id: String,
+    pub update: GooseSessionUpdate,
+}
+
+/// Discriminated union of goose-specific session update payloads.
+/// Variant tag matches ACP's convention (`sessionUpdate: "<snake_case>"`).
+///
+/// `discriminator.mapping` is what makes TS codegen (`@hey-api/openapi-ts`)
+/// emit the correct snake_case tag value even when this enum has a single
+/// variant. Add a mapping entry per variant.
+#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
+#[serde(tag = "sessionUpdate", rename_all = "snake_case")]
+#[schemars(extend("discriminator" = {
+    "propertyName": "sessionUpdate",
+    "mapping": {
+        "usage_update": "#/$defs/SessionUsageUpdate",
+        "status_message": "#/$defs/StatusMessageUpdate",
+        "interaction_update": "#/$defs/InteractionUpdate"
+    }
+}))]
+pub enum GooseSessionUpdate {
+    UsageUpdate(SessionUsageUpdate),
+    StatusMessage(StatusMessageUpdate),
+    InteractionUpdate(InteractionUpdate),
+}
+
+impl Default for GooseSessionUpdate {
+    fn default() -> Self {
+        GooseSessionUpdate::UsageUpdate(SessionUsageUpdate::default())
+    }
+}
+
+/// Streaming context-window usage update for a session.
+#[derive(Debug, Default, Clone, Serialize, Deserialize, JsonSchema)]
+#[serde(rename_all = "camelCase")]
+pub struct SessionUsageUpdate {
+    pub used: u64,
+    pub context_limit: u64,
+    pub accumulated_input_tokens: u64,
+    pub accumulated_output_tokens: u64,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub accumulated_cost: Option<f64>,
+}
+
+/// Live UI/session status. This is not conversation transcript content, and
+/// should not be persisted or replayed as history.
+#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
+#[serde(rename_all = "camelCase")]
+pub struct StatusMessageUpdate {
+    pub status: StatusMessage,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
+#[serde(tag = "type", rename_all = "snake_case")]
+pub enum StatusMessage {
+    #[serde(rename_all = "camelCase")]
+    Notice { message: String },
+    #[serde(rename_all = "camelCase")]
+    Progress { message: String },
+}
+
+#[derive(Debug, Default, Clone, Serialize, Deserialize, JsonSchema)]
+#[serde(rename_all = "camelCase")]
+pub struct InteractionUpdate {
+    pub interaction: Interaction,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    #[serde(rename = "_meta")]
+    pub meta: Option<serde_json::Value>,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
+#[serde(tag = "type", rename_all = "snake_case")]
+pub enum Interaction {
+    #[serde(rename_all = "camelCase")]
+    Elicitation {
+        id: String,
+        state: InteractionState,
+        #[serde(skip_serializing_if = "Option::is_none")]
+        message: Option<String>,
+        #[serde(skip_serializing_if = "Option::is_none")]
+        requested_schema: Option<serde_json::Value>,
+    },
+}
+
+impl Default for Interaction {
+    fn default() -> Self {
+        Self::Elicitation {
+            id: String::new(),
+            state: InteractionState::Pending,
+            message: None,
+            requested_schema: None,
+        }
+    }
+}
+
+#[derive(Debug, Default, Clone, Serialize, Deserialize, JsonSchema)]
+#[serde(rename_all = "snake_case")]
+pub enum InteractionState {
+    #[default]
+    Pending,
+    Submitted,
+}
+
+fn notification_schema<T>(generator: &mut SchemaGenerator) -> CustomMethodSchema
+where
+    T: Default + JsonRpcMessage + JsonSchema,
+{
+    let dummy = T::default();
+    let type_name = std::any::type_name::<T>()
+        .rsplit("::")
+        .next()
+        .unwrap_or(std::any::type_name::<T>())
+        .to_string();
+    CustomMethodSchema {
+        method: dummy.method().to_string(),
+        params_schema: Some(generator.subschema_for::<T>()),
+        params_type_name: Some(type_name),
+        response_schema: None,
+        response_type_name: None,
+    }
+}
+
+/// Schemas for every goose-custom outbound notification. To register a new
+/// notification, define the struct above (with `JsonRpcNotification` +
+/// `Default`) and add one line below.
+pub fn custom_notification_schemas(generator: &mut SchemaGenerator) -> Vec<CustomMethodSchema> {
+    vec![notification_schema::<GooseSessionNotification>(generator)]
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use serde_json::json;
+
+    #[test]
+    fn status_message_serializes_to_expected_wire_shape() {
+        let notification = GooseSessionNotification {
+            session_id: "s1".to_string(),
+            update: GooseSessionUpdate::StatusMessage(StatusMessageUpdate {
+                status: StatusMessage::Notice {
+                    message: "Compaction complete".to_string(),
+                },
+            }),
+        };
+
+        let value = serde_json::to_value(notification).unwrap();
+
+        assert_eq!(
+            value,
+            json!({
+                "sessionId": "s1",
+                "update": {
+                    "sessionUpdate": "status_message",
+                    "status": {
+                        "type": "notice",
+                        "message": "Compaction complete"
+                    }
+                }
+            })
+        );
+    }
+}

crates/goose-sdk/src/custom_requests.rs

@@ -436,6 +436,17 @@ pub struct ImportSessionResponse {
     pub message_count: u64,
 }
 
+/// Submit a response for a pending MCP elicitation in an active session.
+#[derive(Debug, Default, Clone, Serialize, Deserialize, JsonSchema, JsonRpcRequest)]
+#[request(method = "_goose/unstable/elicitation/respond", response = EmptyResponse)]
+#[serde(rename_all = "camelCase")]
+pub struct ElicitationRespondRequest {
+    pub session_id: String,
+    pub elicitation_id: String,
+    #[serde(default)]
+    pub user_data: serde_json::Value,
+}
+
 #[derive(Debug, Default, Clone, Serialize, Deserialize, JsonSchema)]
 #[serde(rename_all = "camelCase")]
 pub struct ProviderConfigKey {

crates/goose-sdk/src/lib.rs

@@ -1 +1,2 @@
+pub mod custom_notifications;
 pub mod custom_requests;