forecite 0.1.3

Official Rust SDK for the Forecite API — scored news feed (REST), Verdict scoring, and realtime WebSocket streaming.
Documentation
use futures_util::{SinkExt, StreamExt};
use serde::Deserialize;
use serde_json::Value;
use tokio::net::TcpStream;
use tokio_tungstenite::{
    connect_async, tungstenite::Message, MaybeTlsStream, WebSocketStream,
};

use crate::error::ForeciteError;
use crate::types::{FeedItem, StreamFilters};

type Ws = WebSocketStream<MaybeTlsStream<TcpStream>>;

/// An event received on the realtime stream.
#[derive(Debug, Clone, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum StreamEvent {
    Welcome {
        tier: String,
        limits: Value,
    },
    Feed {
        #[serde(default)]
        snapshot: bool,
        #[serde(default)]
        count: Option<u64>,
        data: FeedItem,
    },
    QuotaExceeded {
        limit: i64,
        message: String,
    },
    Error {
        code: Option<String>,
        message: Option<String>,
    },
    /// `subscribed`, `pong`, or anything else we don't model.
    #[serde(other)]
    Other,
}

/// A live feed subscription. Drive it by calling [`ForeciteStream::next`] in a
/// loop; it returns `Ok(None)` when the connection closes.
pub struct ForeciteStream {
    ws: Ws,
}

impl ForeciteStream {
    pub(crate) async fn connect(
        ws_url: &str,
        api_key: &str,
        filters: &StreamFilters,
        snapshot: u32,
    ) -> Result<Self, ForeciteError> {
        let url = format!("{ws_url}/?api_key={}", urlencoding::encode(api_key));
        let (mut ws, _) = connect_async(url).await?;
        let subscribe = serde_json::json!({
            "type": "subscribe",
            "filters": filters,
            "snapshot": snapshot,
        });
        ws.send(Message::Text(subscribe.to_string().into())).await?;
        Ok(Self { ws })
    }

    /// Receive the next event. Auto-responds to pings. `Ok(None)` = closed.
    pub async fn next(&mut self) -> Result<Option<StreamEvent>, ForeciteError> {
        while let Some(message) = self.ws.next().await {
            match message? {
                Message::Text(text) => {
                    return Ok(Some(serde_json::from_str(&text)?));
                }
                Message::Ping(payload) => {
                    self.ws.send(Message::Pong(payload)).await?;
                }
                Message::Close(_) => return Ok(None),
                _ => {}
            }
        }
        Ok(None)
    }

    /// Close the connection.
    pub async fn close(mut self) -> Result<(), ForeciteError> {
        self.ws.close(None).await?;
        Ok(())
    }
}