systemprompt-agent 0.2.0

Core Agent protocol module for systemprompt.io
Documentation
use systemprompt_identifiers::MessageId;
use systemprompt_models::{AgUiEventBuilder, AgUiMessageRole};

use crate::services::a2a_server::streaming::webhook_client::WebhookContext;

pub struct TextStreamState {
    message_started: bool,
    webhook_context: Option<WebhookContext>,
}

impl TextStreamState {
    pub const fn new() -> Self {
        Self {
            message_started: false,
            webhook_context: None,
        }
    }

    pub fn with_webhook_context(mut self, context: WebhookContext) -> Self {
        self.webhook_context = Some(context);
        self
    }

    pub async fn handle_text(&mut self, text: String, message_id: &MessageId) {
        let Some(ref webhook_context) = self.webhook_context else {
            return;
        };

        let message_id_str = message_id.as_str();

        if !self.message_started {
            let start_event =
                AgUiEventBuilder::text_message_start(message_id_str, AgUiMessageRole::Assistant);
            if let Err(e) = webhook_context.broadcast_agui(start_event).await {
                tracing::error!(error = %e, "Failed to broadcast TEXT_MESSAGE_START");
            }
            self.message_started = true;
        }

        let content_event = AgUiEventBuilder::text_message_content(message_id_str, &text);
        if let Err(e) = webhook_context.broadcast_agui(content_event).await {
            tracing::error!(error = %e, "Failed to broadcast TEXT_MESSAGE_CONTENT");
        }
    }

    pub async fn finalize(&self, message_id: &MessageId) {
        if self.message_started {
            if let Some(ref webhook_context) = self.webhook_context {
                let end_event = AgUiEventBuilder::text_message_end(message_id.as_str());
                if let Err(e) = webhook_context.broadcast_agui(end_event).await {
                    tracing::error!(error = %e, "Failed to broadcast TEXT_MESSAGE_END");
                }
            }
        }
    }
}

impl Default for TextStreamState {
    fn default() -> Self {
        Self::new()
    }
}