Skip to main content

tea_protocol/
content.rs

1use base64::Engine as _;
2use base64::engine::general_purpose::STANDARD;
3use serde::{Deserialize, Deserializer, Serialize, Serializer};
4use serde_json::Value;
5use thiserror::Error;
6
7use crate::ToolCallId;
8use crate::external::{ExternalContentError, HostedToolActivity, SourceCitation};
9use crate::metadata::{ProtocolMetadataError, validate_json_bounds};
10
11/// Maximum UTF-8 bytes in one text or thinking content block.
12pub const MAX_TEXT_BLOCK_BYTES: usize = 1024 * 1024;
13/// Maximum encoded bytes in one inline Base64 image.
14pub const MAX_INLINE_IMAGE_BASE64_BYTES: usize = 6 * 1024 * 1024;
15/// Maximum encoded JSON bytes for tool arguments.
16pub const MAX_TOOL_ARGUMENT_BYTES: usize = 256 * 1024;
17/// Maximum nesting depth for tool arguments.
18pub const MAX_TOOL_ARGUMENT_DEPTH: usize = 32;
19/// Maximum UTF-8 bytes in one opaque provider tool-call identifier.
20pub const MAX_PROVIDER_TOOL_CALL_ID_BYTES: usize = 256;
21
22/// A provider-neutral message content block.
23#[derive(Debug, Clone, PartialEq)]
24pub enum ContentBlock {
25    /// Visible text content.
26    Text {
27        /// UTF-8 text.
28        text: String,
29    },
30    /// Model reasoning content that products may choose to hide.
31    Thinking {
32        /// UTF-8 reasoning text.
33        text: String,
34    },
35    /// Image content with a MIME type and source.
36    Image {
37        /// Image MIME type.
38        mime_type: String,
39        /// Inline or referenced image source.
40        source: ImageSource,
41    },
42    /// A request to invoke a named tool.
43    ToolCall {
44        /// Canonical tool-call identifier.
45        tool_call_id: ToolCallId,
46        /// Opaque provider identifier used to continue the model conversation.
47        provider_call_id: Option<String>,
48        /// Registered tool name.
49        tool_name: String,
50        /// Provider-neutral JSON arguments.
51        arguments: Value,
52    },
53    /// One complete provider-hosted activity, including normalized sources.
54    HostedTool {
55        /// Validated activity and opaque same-provider continuation state.
56        activity: HostedToolActivity,
57    },
58    /// A normalized citation associated with assistant text.
59    Citation {
60        /// Validated cited source and optional provider continuation state.
61        citation: SourceCitation,
62    },
63}
64
65impl ContentBlock {
66    /// Creates a validated visible text block.
67    ///
68    /// # Errors
69    ///
70    /// Returns [`ContentValidationError::InvalidText`] when the text is empty,
71    /// contains a null character, or exceeds [`MAX_TEXT_BLOCK_BYTES`].
72    pub fn text(text: impl Into<String>) -> Result<Self, ContentValidationError> {
73        let text = text.into();
74        validate_text(&text)?;
75        Ok(Self::Text { text })
76    }
77
78    /// Creates a validated thinking block.
79    ///
80    /// # Errors
81    ///
82    /// Returns [`ContentValidationError::InvalidText`] when the text is empty,
83    /// contains a null character, or exceeds [`MAX_TEXT_BLOCK_BYTES`].
84    pub fn thinking(text: impl Into<String>) -> Result<Self, ContentValidationError> {
85        let text = text.into();
86        validate_text(&text)?;
87        Ok(Self::Thinking { text })
88    }
89
90    /// Creates a validated inline Base64 image block.
91    ///
92    /// # Errors
93    ///
94    /// Returns an error when the MIME type is invalid or the data is empty,
95    /// oversized, or not valid standard Base64.
96    pub fn inline_image(
97        mime_type: impl Into<String>,
98        data: impl Into<String>,
99    ) -> Result<Self, ContentValidationError> {
100        let mime_type = mime_type.into();
101        let data = data.into();
102        validate_mime_type(&mime_type)?;
103        if data.is_empty() || data.len() > MAX_INLINE_IMAGE_BASE64_BYTES {
104            return Err(ContentValidationError::InvalidImageData);
105        }
106        STANDARD
107            .decode(data.as_bytes())
108            .map_err(|_| ContentValidationError::InvalidImageData)?;
109        Ok(Self::Image {
110            mime_type,
111            source: ImageSource::InlineBase64 { data },
112        })
113    }
114
115    /// Creates a validated referenced image block.
116    ///
117    /// # Errors
118    ///
119    /// Returns an error when the MIME type or bounded reference is invalid.
120    pub fn image_reference(
121        mime_type: impl Into<String>,
122        reference: impl Into<String>,
123    ) -> Result<Self, ContentValidationError> {
124        let mime_type = mime_type.into();
125        let reference = reference.into();
126        validate_mime_type(&mime_type)?;
127        if reference.is_empty() || reference.len() > 1024 || reference.chars().any(char::is_control)
128        {
129            return Err(ContentValidationError::InvalidImageReference);
130        }
131        Ok(Self::Image {
132            mime_type,
133            source: ImageSource::Reference { reference },
134        })
135    }
136
137    /// Creates a validated tool-call block.
138    ///
139    /// # Errors
140    ///
141    /// Returns an error when the tool name is invalid, arguments are not an
142    /// object, or arguments exceed the byte or nesting limits.
143    pub fn tool_call(
144        tool_call_id: ToolCallId,
145        tool_name: impl Into<String>,
146        arguments: Value,
147    ) -> Result<Self, ContentValidationError> {
148        Self::tool_call_inner(tool_call_id, None, tool_name.into(), arguments)
149    }
150
151    /// Creates a validated tool-call block that retains its provider identifier.
152    ///
153    /// # Errors
154    ///
155    /// Returns an error when the provider identifier, tool name, or arguments
156    /// violate their protocol bounds.
157    pub fn tool_call_with_provider_id(
158        tool_call_id: ToolCallId,
159        provider_call_id: impl Into<String>,
160        tool_name: impl Into<String>,
161        arguments: Value,
162    ) -> Result<Self, ContentValidationError> {
163        Self::tool_call_inner(
164            tool_call_id,
165            Some(provider_call_id.into()),
166            tool_name.into(),
167            arguments,
168        )
169    }
170
171    fn tool_call_inner(
172        tool_call_id: ToolCallId,
173        provider_call_id: Option<String>,
174        tool_name: String,
175        arguments: Value,
176    ) -> Result<Self, ContentValidationError> {
177        if let Some(provider_call_id) = provider_call_id.as_deref() {
178            validate_provider_tool_call_id(provider_call_id)?;
179        }
180        validate_tool_name(&tool_name)?;
181        if !arguments.is_object() {
182            return Err(ContentValidationError::ToolArgumentsMustBeObject);
183        }
184        validate_json_bounds(&arguments, MAX_TOOL_ARGUMENT_BYTES, MAX_TOOL_ARGUMENT_DEPTH)?;
185        Ok(Self::ToolCall {
186            tool_call_id,
187            provider_call_id,
188            tool_name,
189            arguments,
190        })
191    }
192
193    /// Wraps one validated provider-hosted tool activity.
194    #[must_use]
195    pub fn hosted_tool(activity: HostedToolActivity) -> Self {
196        Self::HostedTool { activity }
197    }
198
199    /// Wraps one validated external source citation.
200    #[must_use]
201    pub fn citation(citation: SourceCitation) -> Self {
202        Self::Citation { citation }
203    }
204
205    /// Returns the opaque provider identifier for a tool-call block.
206    #[must_use]
207    pub fn provider_call_id(&self) -> Option<&str> {
208        match self {
209            Self::ToolCall {
210                provider_call_id, ..
211            } => provider_call_id.as_deref(),
212            Self::HostedTool { activity } => Some(activity.provider_call_id()),
213            _ => None,
214        }
215    }
216
217    pub(crate) fn validate(&self) -> Result<(), ContentValidationError> {
218        match self {
219            Self::Text { text } | Self::Thinking { text } => validate_text(text),
220            Self::Image { mime_type, source } => {
221                validate_mime_type(mime_type)?;
222                match source {
223                    ImageSource::InlineBase64 { data } => {
224                        if data.is_empty() || data.len() > MAX_INLINE_IMAGE_BASE64_BYTES {
225                            return Err(ContentValidationError::InvalidImageData);
226                        }
227                        STANDARD
228                            .decode(data.as_bytes())
229                            .map_err(|_| ContentValidationError::InvalidImageData)?;
230                    }
231                    ImageSource::Reference { reference } => {
232                        if reference.is_empty()
233                            || reference.len() > 1024
234                            || reference.chars().any(char::is_control)
235                        {
236                            return Err(ContentValidationError::InvalidImageReference);
237                        }
238                    }
239                }
240                Ok(())
241            }
242            Self::ToolCall {
243                provider_call_id,
244                tool_name,
245                arguments,
246                ..
247            } => {
248                if let Some(provider_call_id) = provider_call_id.as_deref() {
249                    validate_provider_tool_call_id(provider_call_id)?;
250                }
251                validate_tool_name(tool_name)?;
252                if !arguments.is_object() {
253                    return Err(ContentValidationError::ToolArgumentsMustBeObject);
254                }
255                validate_json_bounds(arguments, MAX_TOOL_ARGUMENT_BYTES, MAX_TOOL_ARGUMENT_DEPTH)?;
256                Ok(())
257            }
258            Self::HostedTool { activity } => activity.validate().map_err(Into::into),
259            Self::Citation { citation } => citation.validate().map_err(Into::into),
260        }
261    }
262
263    pub(crate) const fn valid_for_user(&self) -> bool {
264        matches!(self, Self::Text { .. } | Self::Image { .. })
265    }
266
267    pub(crate) const fn valid_for_assistant(&self) -> bool {
268        matches!(
269            self,
270            Self::Text { .. }
271                | Self::Thinking { .. }
272                | Self::ToolCall { .. }
273                | Self::HostedTool { .. }
274                | Self::Citation { .. }
275        )
276    }
277
278    pub(crate) const fn valid_for_tool_result(&self) -> bool {
279        matches!(self, Self::Text { .. } | Self::Image { .. })
280    }
281}
282
283impl Serialize for ContentBlock {
284    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
285    where
286        S: Serializer,
287    {
288        self.validate().map_err(serde::ser::Error::custom)?;
289        SerializableContentBlock::from(self).serialize(serializer)
290    }
291}
292
293#[derive(Serialize)]
294#[serde(tag = "type", rename_all = "snake_case")]
295enum SerializableContentBlock<'a> {
296    Text {
297        text: &'a str,
298    },
299    Thinking {
300        text: &'a str,
301    },
302    Image {
303        #[serde(rename = "mimeType")]
304        mime_type: &'a str,
305        source: &'a ImageSource,
306    },
307    ToolCall {
308        #[serde(rename = "toolCallId")]
309        tool_call_id: &'a ToolCallId,
310        #[serde(rename = "providerCallId", skip_serializing_if = "Option::is_none")]
311        provider_call_id: Option<&'a str>,
312        #[serde(rename = "toolName")]
313        tool_name: &'a str,
314        arguments: &'a Value,
315    },
316    HostedTool {
317        activity: &'a HostedToolActivity,
318    },
319    Citation {
320        citation: &'a SourceCitation,
321    },
322}
323
324impl<'a> From<&'a ContentBlock> for SerializableContentBlock<'a> {
325    fn from(value: &'a ContentBlock) -> Self {
326        match value {
327            ContentBlock::Text { text } => Self::Text { text },
328            ContentBlock::Thinking { text } => Self::Thinking { text },
329            ContentBlock::Image { mime_type, source } => Self::Image { mime_type, source },
330            ContentBlock::ToolCall {
331                tool_call_id,
332                provider_call_id,
333                tool_name,
334                arguments,
335            } => Self::ToolCall {
336                tool_call_id,
337                provider_call_id: provider_call_id.as_deref(),
338                tool_name,
339                arguments,
340            },
341            ContentBlock::HostedTool { activity } => Self::HostedTool { activity },
342            ContentBlock::Citation { citation } => Self::Citation { citation },
343        }
344    }
345}
346
347#[derive(Deserialize)]
348#[serde(tag = "type", rename_all = "snake_case")]
349enum RawContentBlock {
350    Text {
351        text: String,
352    },
353    Thinking {
354        text: String,
355    },
356    Image {
357        #[serde(rename = "mimeType")]
358        mime_type: String,
359        source: ImageSource,
360    },
361    ToolCall {
362        #[serde(rename = "toolCallId")]
363        tool_call_id: ToolCallId,
364        #[serde(rename = "providerCallId", default)]
365        provider_call_id: Option<String>,
366        #[serde(rename = "toolName")]
367        tool_name: String,
368        arguments: Value,
369    },
370    HostedTool {
371        activity: HostedToolActivity,
372    },
373    Citation {
374        citation: SourceCitation,
375    },
376}
377
378impl<'de> Deserialize<'de> for ContentBlock {
379    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
380    where
381        D: Deserializer<'de>,
382    {
383        let raw = RawContentBlock::deserialize(deserializer)?;
384        let result = match raw {
385            RawContentBlock::Text { text } => Self::text(text),
386            RawContentBlock::Thinking { text } => Self::thinking(text),
387            RawContentBlock::Image { mime_type, source } => match source {
388                ImageSource::InlineBase64 { data } => Self::inline_image(mime_type, data),
389                ImageSource::Reference { reference } => Self::image_reference(mime_type, reference),
390            },
391            RawContentBlock::ToolCall {
392                tool_call_id,
393                provider_call_id,
394                tool_name,
395                arguments,
396            } => Self::tool_call_inner(tool_call_id, provider_call_id, tool_name, arguments),
397            RawContentBlock::HostedTool { activity } => activity
398                .validate()
399                .map_err(ContentValidationError::from)
400                .map(|()| Self::hosted_tool(activity)),
401            RawContentBlock::Citation { citation } => citation
402                .validate()
403                .map_err(ContentValidationError::from)
404                .map(|()| Self::citation(citation)),
405        };
406        result.map_err(serde::de::Error::custom)
407    }
408}
409
410/// Source of an image content block.
411#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
412#[serde(tag = "type", rename_all = "snake_case")]
413pub enum ImageSource {
414    /// Base64 data embedded directly in the protocol value.
415    InlineBase64 {
416        /// Standard padded Base64 text.
417        data: String,
418    },
419    /// Stable reference resolved by an adapter or artifact store.
420    Reference {
421        /// Bounded opaque reference string.
422        reference: String,
423    },
424}
425
426/// Error returned when validating content blocks.
427#[derive(Debug, Error)]
428pub enum ContentValidationError {
429    /// Text is empty, too large, or contains a null character.
430    #[error("text content is empty, too large, or contains a null character")]
431    InvalidText,
432    /// The MIME type is not a supported canonical image type.
433    #[error("image MIME type must use canonical image/type syntax")]
434    InvalidMimeType,
435    /// Inline image data is empty, too large, or not valid standard Base64.
436    #[error("inline image data is invalid")]
437    InvalidImageData,
438    /// An image reference is empty, too large, or contains controls.
439    #[error("image reference is invalid")]
440    InvalidImageReference,
441    /// The tool name is not canonical.
442    #[error(
443        "tool name must start with a lowercase letter and contain lowercase ASCII, digits, '_', '-', or '.'"
444    )]
445    InvalidToolName,
446    /// The provider tool-call identifier is empty, oversized, or contains controls.
447    #[error("provider tool-call identifier is invalid")]
448    InvalidProviderToolCallId,
449    /// Tool arguments must be a JSON object.
450    #[error("tool arguments must be a JSON object")]
451    ToolArgumentsMustBeObject,
452    /// Tool arguments exceed JSON byte or nesting limits.
453    #[error("tool arguments exceed protocol bounds: {0}")]
454    ToolArgumentsOutOfBounds(#[from] ProtocolMetadataError),
455    /// Hosted activity, source, citation, or continuation content is invalid.
456    #[error("external content is invalid: {0}")]
457    InvalidExternalContent(#[from] ExternalContentError),
458}
459
460fn validate_text(text: &str) -> Result<(), ContentValidationError> {
461    if text.is_empty() || text.len() > MAX_TEXT_BLOCK_BYTES || text.contains('\0') {
462        Err(ContentValidationError::InvalidText)
463    } else {
464        Ok(())
465    }
466}
467
468fn validate_mime_type(value: &str) -> Result<(), ContentValidationError> {
469    let subtype = value
470        .strip_prefix("image/")
471        .ok_or(ContentValidationError::InvalidMimeType)?;
472    if subtype.is_empty()
473        || !subtype.bytes().all(|byte| {
474            byte.is_ascii_lowercase() || byte.is_ascii_digit() || matches!(byte, b'+' | b'-' | b'.')
475        })
476    {
477        return Err(ContentValidationError::InvalidMimeType);
478    }
479    Ok(())
480}
481
482pub(crate) fn validate_tool_name(value: &str) -> Result<(), ContentValidationError> {
483    let mut bytes = value.bytes();
484    if !bytes.next().is_some_and(|byte| byte.is_ascii_lowercase())
485        || value.len() > 128
486        || !bytes.all(|byte| {
487            byte.is_ascii_lowercase() || byte.is_ascii_digit() || matches!(byte, b'_' | b'-' | b'.')
488        })
489    {
490        return Err(ContentValidationError::InvalidToolName);
491    }
492    Ok(())
493}
494
495pub(crate) fn validate_provider_tool_call_id(value: &str) -> Result<(), ContentValidationError> {
496    if value.is_empty()
497        || value.len() > MAX_PROVIDER_TOOL_CALL_ID_BYTES
498        || value.chars().any(char::is_control)
499    {
500        Err(ContentValidationError::InvalidProviderToolCallId)
501    } else {
502        Ok(())
503    }
504}