//! System message transform callbacks for customizing agent prompts.

//!

//! Implement [`SystemMessageTransform`](crate::transforms::SystemMessageTransform) to intercept and modify system prompt

//! sections during session creation. The CLI sends the current content for

//! each section the transform registered, and the SDK returns the modified

//! content.

use std::collections::HashMap;

use async_trait::async_trait;

use serde::{Deserialize, Serialize};

use crate::types::SessionId;

/// Context provided to every transform invocation.

#[derive(Debug, Clone)]

pub struct TransformContext {

/// The session being created or resumed.

pub session_id: SessionId,

}

/// Handles `systemMessage.transform` RPC requests from the CLI.

///

/// The CLI sends these during session creation/resumption when the session's

/// `SystemMessageConfig` contains sections with `action: "transform"`. For each

/// such section, the CLI provides the current content and expects the SDK to

/// return the (possibly modified) content.

///

/// Implement this trait and pass it to [`Client::create_session`](crate::Client::create_session) /

/// [`Client::resume_session`](crate::Client::resume_session) to participate in system message customization.

///

/// # Example

///

/// ```ignore

/// struct MyTransform;

///

/// #[async_trait::async_trait]

/// impl SystemMessageTransform for MyTransform {

/// fn section_ids(&self) -> Vec<String> {

/// vec!["instructions".to_string()]

/// }

///

/// async fn transform_section(

/// &self,

/// _section_id: &str,

/// content: &str,

/// _ctx: TransformContext,

/// ) -> Option<String> {

/// Some(format!("{content}\n\nAlways be concise."))

/// }

/// }

/// ```

#[async_trait]

pub trait SystemMessageTransform: Send + Sync + 'static {

/// Section IDs this transform handles.

///

/// The SDK injects `action: "transform"` entries into the

/// [`SystemMessageConfig`](crate::types::SystemMessageConfig) wire format

/// for each returned ID.

fn section_ids(&self) -> Vec<String>;

/// Transform a section's content. Return `Some(new_content)` to modify the

/// section, or `None` to pass through unchanged.

async fn transform_section(

&self,

section_id: &str,

content: &str,

ctx: TransformContext,

) -> Option<String>;

}

/// Wire format for a single section in the transform request/response.

#[derive(Debug, Clone, Serialize, Deserialize)]

pub(crate) struct TransformSection {

pub(crate) content: String,

}

/// Wire format for the `systemMessage.transform` response.

#[derive(Debug, Clone, Serialize)]

pub(crate) struct TransformResponse {

pub(crate) sections: HashMap<String, TransformSection>,

}

/// Apply transforms to the incoming sections map, returning the response.

///

/// For each section, calls the matching transform if the implementor returns

/// `Some`; otherwise passes through the original content.

pub(crate) async fn dispatch_transform(

transform: &dyn SystemMessageTransform,

session_id: &SessionId,

sections: HashMap<String, TransformSection>,

) -> TransformResponse {

let ctx = TransformContext {

session_id: session_id.clone(),

};

let mut result = HashMap::with_capacity(sections.len());

for (section_id, data) in sections {

let content = match transform

.transform_section(&section_id, &data.content, ctx.clone())

.await

{

Some(transformed) => transformed,

None => data.content,

};

result.insert(section_id, TransformSection { content });

}

TransformResponse { sections: result }

}

#[cfg(test)]

mod tests {

use super::*;

struct TestTransform;

#[async_trait]

impl SystemMessageTransform for TestTransform {

fn section_ids(&self) -> Vec<String> {

vec!["instructions".to_string(), "context".to_string()]

}

async fn transform_section(

&self,

section_id: &str,

content: &str,

_ctx: TransformContext,

) -> Option<String> {

match section_id {

"instructions" => Some(format!("[modified] {content}")),

_ => None,

}

}

}

#[tokio::test]

async fn dispatch_applies_matching_transform() {

let transform = TestTransform;

let mut sections = HashMap::new();

sections.insert(

"instructions".to_string(),

TransformSection {

content: "be helpful".to_string(),

},

);

let response = dispatch_transform(&transform, &SessionId::new("sess-1"), sections).await;

assert_eq!(

response.sections["instructions"].content,

"[modified] be helpful"

);

}

#[tokio::test]

async fn dispatch_passes_through_unhandled_section() {

let transform = TestTransform;

let mut sections = HashMap::new();

sections.insert(

"context".to_string(),

TransformSection {

content: "original context".to_string(),

},

);

let response = dispatch_transform(&transform, &SessionId::new("sess-1"), sections).await;

assert_eq!(response.sections["context"].content, "original context");

}

#[tokio::test]

async fn dispatch_unknown_section_passes_through() {

let transform = TestTransform;

let mut sections = HashMap::new();

sections.insert(

"unknown".to_string(),

TransformSection {

content: "mystery".to_string(),

},

);

let response = dispatch_transform(&transform, &SessionId::new("sess-1"), sections).await;

assert_eq!(response.sections["unknown"].content, "mystery");

}

#[tokio::test]

async fn dispatch_mixed_sections() {

let transform = TestTransform;

let mut sections = HashMap::new();

sections.insert(

"instructions".to_string(),

TransformSection {

content: "help me".to_string(),

},

);

sections.insert(

"context".to_string(),

TransformSection {

content: "some context".to_string(),

},

);

sections.insert(

"other".to_string(),

TransformSection {

content: "other stuff".to_string(),

},

);

let response = dispatch_transform(&transform, &SessionId::new("sess-1"), sections).await;

assert_eq!(

response.sections["instructions"].content,

"[modified] help me"

);

assert_eq!(response.sections["context"].content, "some context");

assert_eq!(response.sections["other"].content, "other stuff");

}

#[tokio::test]

async fn section_ids_returns_registered_sections() {

let transform = TestTransform;

let ids = transform.section_ids();

assert_eq!(ids, vec!["instructions", "context"]);

}

}