Something went wrong. Try again.
This repository has no description
Something went wrong. Try again.
9.2 kB · 287 lines
Rust
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287use anyhow::{Context, Result};use futures_util::StreamExt;use serde::{Deserialize, Serialize};use std::{collections::HashMap, time::Duration};use tokio::sync::mpsc;use tokio_tungstenite::{ connect_async, tungstenite::{ client::IntoClientRequest, http::HeaderValue, Message, },};use url::Url;
use crate::tui::Message as TuiMessage;
#[derive(Debug, Clone, Serialize, Deserialize)]pub struct JetstreamEvent { #[serde(rename = "kind")] pub kind: String, #[serde(rename = "time_us")] pub time_us: i64, pub did: String, pub commit: Option<CommitData>,}
#[derive(Debug, Clone, Serialize, Deserialize)]pub struct CommitData { pub rev: String, pub operation: String, pub collection: String, pub rkey: String, pub record: Option<serde_json::Value>, pub cid: String,}
#[derive(Debug, Clone, Serialize, Deserialize)]pub struct BlipRecord { #[serde(rename = "$type")] pub record_type: String, pub content: String, #[serde(rename = "createdAt")] pub created_at: String,}
pub struct JetstreamClient { did_cache: HashMap<String, String>, // DID -> handle cache own_did: Option<String>, // User's own DID to filter out}
impl JetstreamClient { pub fn new(own_did: Option<String>) -> Self { Self { did_cache: HashMap::new(), own_did, } }
pub async fn connect_and_listen(&mut self, message_tx: mpsc::UnboundedSender<TuiMessage>) -> Result<()> { // Try simple connection first, then with collection filter let urls = vec![ "wss://jetstream2.us-west.bsky.network/subscribe", "wss://jetstream2.us-west.bsky.network/subscribe?wantedCollections=stream.thought.blip" ]; for (i, jetstream_url) in urls.iter().enumerate() { // Send status to TUI instead of console let status_msg = crate::tui::Message::new( "system".to_string(), format!("Trying connection {} of {}", i + 1, urls.len()), false, ); let _ = message_tx.send(status_msg); loop { match self.try_connect_and_listen(&message_tx, jetstream_url).await { Ok(_) => { return Ok(()); } Err(_e) => { // If this is the last URL, retry it after a delay if i == urls.len() - 1 { let retry_msg = crate::tui::Message::new( "system".to_string(), "Connection failed, retrying in 5s...".to_string(), false, ); let _ = message_tx.send(retry_msg); tokio::time::sleep(Duration::from_secs(5)).await; } else { // Try the next URL break; } } } } } Ok(()) }
async fn try_connect_and_listen( &mut self, message_tx: &mpsc::UnboundedSender<TuiMessage>, url_str: &str, ) -> Result<()> { // Parse URL and create request with headers let url = Url::parse(url_str)?; let mut request = url.into_client_request()?; // Add User-Agent header request.headers_mut().insert( "User-Agent", HeaderValue::from_static("think-cli/0.1.0") ); // Connect with timeout let connect_future = connect_async(request); let (ws_stream, _response) = tokio::time::timeout( Duration::from_secs(10), connect_future ).await .context("Connection timeout")? .context("Failed to connect to jetstream")?;
// Send a connection success message to the TUI let success_msg = crate::tui::Message::new( "system".to_string(), "Connected to jetstream! Listening for blips...".to_string(), false, ); let _ = message_tx.send(success_msg);
let (mut _write, mut read) = ws_stream.split();
while let Some(msg) = read.next().await { match msg { Ok(Message::Text(text)) => { // Silently ignore message handling errors let _ = self.handle_message(&text, message_tx).await; } Ok(Message::Close(_)) => { break; } Err(e) => { return Err(anyhow::anyhow!("WebSocket error: {}", e)); } _ => { // Ignore other message types (binary, ping, pong) } } }
Ok(()) }
async fn handle_message( &mut self, message: &str, message_tx: &mpsc::UnboundedSender<TuiMessage>, ) -> Result<()> { // First, check if it's even a commit event using basic JSON parsing let event_value: serde_json::Value = serde_json::from_str(message)?; // Only process commit events if event_value.get("kind").and_then(|k| k.as_str()) != Some("commit") { return Ok(()); } // Check if it has a commit with the right collection let commit = event_value.get("commit"); if let Some(commit_obj) = commit { if commit_obj.get("collection").and_then(|c| c.as_str()) != Some("stream.thought.blip") { return Ok(()); } // Skip delete operations if commit_obj.get("operation").and_then(|o| o.as_str()) == Some("delete") { return Ok(()); } } else { return Ok(()); }
// Now try to parse as our structured event let event: JetstreamEvent = serde_json::from_str(message)?; let commit = event.commit.as_ref().unwrap(); // Safe because we checked above
// Skip messages from our own DID if let Some(ref own_did) = self.own_did { if &event.did == own_did { return Ok(()); } }
// Parse the blip record let record_data = commit.record.as_ref(); if record_data.is_none() { return Ok(()); }
let blip_record: BlipRecord = match serde_json::from_value(record_data.unwrap().clone()) { Ok(record) => record, Err(_) => return Ok(()), // Silently skip unparseable records };
// Get or resolve the handle let handle = self.resolve_did(&event.did).await;
// Create TUI message let tui_message = TuiMessage::new( handle, blip_record.content, false, // Not our own message );
// Send to TUI if let Err(_) = message_tx.send(tui_message) { // Channel closed, probably shutting down return Err(anyhow::anyhow!("Message channel closed")); }
Ok(()) }
async fn resolve_did(&mut self, did: &str) -> String { // Check cache first if let Some(handle) = self.did_cache.get(did) { return handle.clone(); }
// Try to resolve the DID to a handle let handle = match self.fetch_handle_for_did(did).await { Ok(h) => h, Err(_) => { // Fallback to showing just the DID (truncated) if did.len() > 20 { format!("{}...", &did[..20]) } else { did.to_string() } } };
// Cache the result self.did_cache.insert(did.to_string(), handle.clone());
handle }
async fn fetch_handle_for_did(&self, did: &str) -> Result<String> { // Use the ATProto API to resolve DID to handle let client = reqwest::Client::new(); let url = format!("https://public.api.bsky.app/xrpc/com.atproto.identity.resolveHandle?did={}", did); #[derive(Deserialize)] struct ResolveResponse { handle: String, }
// Try a simpler approach - resolve via profile let profile_url = format!("https://public.api.bsky.app/xrpc/app.bsky.actor.getProfile?actor={}", did); #[derive(Deserialize)] struct ProfileResponse { handle: String, }
let response = client .get(&profile_url) .send() .await?;
if response.status().is_success() { let profile: ProfileResponse = response.json().await?; Ok(profile.handle) } else { Err(anyhow::anyhow!("Failed to resolve DID to handle")) } }}
pub async fn start_jetstream_listener(message_tx: mpsc::UnboundedSender<TuiMessage>, own_did: Option<String>) -> Result<()> { let mut client = JetstreamClient::new(own_did); client.connect_and_listen(message_tx).await}