Skip to main content

qubit_http/sse/
done_marker_policy.rs

1/*******************************************************************************
2 *
3 *    Copyright (c) 2025 - 2026 Haixing Hu.
4 *
5 *    SPDX-License-Identifier: Apache-2.0
6 *
7 *    Licensed under the Apache License, Version 2.0.
8 *
9 ******************************************************************************/
10//! # Done Marker Policy
11//!
12//! Defines how stream completion markers are recognized.
13//!
14
15use serde::{
16    Deserialize,
17    Serialize,
18};
19use std::fmt;
20use std::str::FromStr;
21
22/// Policy for stream completion marker matching.
23#[derive(Debug, Clone, PartialEq, Eq, Default, Serialize, Deserialize)]
24#[serde(rename_all = "snake_case")]
25pub enum DoneMarkerPolicy {
26    /// Disable done marker recognition.
27    Disabled,
28    /// Use default marker: `[DONE]`.
29    #[default]
30    DefaultDone,
31    /// Use a custom marker string.
32    Custom(String),
33}
34
35impl fmt::Display for DoneMarkerPolicy {
36    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
37        match self {
38            Self::Disabled => write!(f, "disable"),
39            Self::DefaultDone => write!(f, "default"),
40            Self::Custom(value) => write!(f, "{value}"),
41        }
42    }
43}
44
45impl FromStr for DoneMarkerPolicy {
46    type Err = parse_display::ParseError;
47
48    fn from_str(input: &str) -> Result<Self, Self::Err> {
49        let trimmed = input.trim();
50        if trimmed.eq_ignore_ascii_case("disable") {
51            return Ok(Self::Disabled);
52        }
53        if trimmed.eq_ignore_ascii_case("disabled") {
54            return Ok(Self::Disabled);
55        }
56        if trimmed.eq_ignore_ascii_case("default") {
57            return Ok(Self::DefaultDone);
58        }
59        Ok(Self::Custom(trimmed.to_string()))
60    }
61}
62
63impl DoneMarkerPolicy {
64    /// Returns whether `payload` (trimmed) signals end-of-stream per this policy.
65    ///
66    /// # Parameters
67    /// - `payload`: Typically trimmed SSE `data:` text.
68    ///
69    /// # Returns
70    /// `true` when the stream should stop emitting data chunks (e.g. `[DONE]`).
71    pub fn is_done(&self, payload: &str) -> bool {
72        match self {
73            DoneMarkerPolicy::Disabled => false,
74            DoneMarkerPolicy::DefaultDone => payload.trim() == "[DONE]",
75            DoneMarkerPolicy::Custom(marker) => payload.trim() == marker.trim(),
76        }
77    }
78}