openagent 0.1.10

//! HTTP client abstraction and implementations.

// std
use std::{
	env,
	fmt::Debug,
	io::{Error as IoError, Result as IoResult},
	mem,
	pin::Pin,
	task::{Context, Poll},
	time::Duration,
};
// crates.io
use futures::{Stream, TryStreamExt};
use reqwest::{
	Body, Client,
	multipart::{Form, Part},
};
use tokio_util::{
	bytes::Bytes,
	codec::{FramedRead, LinesCodec},
	io::StreamReader,
};
// self
use crate::_prelude::*;

pub(crate) type EventStream<T> = _Stream<Result<T>>;

type _Stream<T> = Pin<Box<dyn Send + Stream<Item = T>>>;
type ByteStream = _Stream<IoResult<Bytes>>;

/// Top-level HTTP abstraction trait defining all client capabilities.
pub trait ApiBase
where
	Self: Send + Sync,
{
	/// Returns the root URL that is prepended to every request path.
	fn base_uri(&self) -> &str;

	/// Issues a GET request and returns the full response body as `String`.
	fn get(&self, endpoint: &str) -> impl Send + Future<Output = Result<String>>;

	/// Issues a multipart POST request and returns the full response body as `String`.
	fn post_multipart(
		&self,
		endpoint: &str,
		multipart: Multipart,
	) -> impl Send + Future<Output = Result<String>>;

	/// Issues a JSON POST request and returns the full response body as `String`.
	fn post_json<S>(&self, endpoint: &str, body: S) -> impl Send + Future<Output = Result<String>>
	where
		S: Send + Serialize;

	/// Performs a streaming POST request and yields server-sent events.
	fn sse<S, H>(
		&self,
		endpoint: &str,
		body: S,
		options: SseOptions<H>,
	) -> impl Send + Future<Output = Result<EventStream<H::Event>>>
	where
		S: Send + Serialize,
		H: 'static + EventHandler;

	/// Same as `sse` but supports resuming from a given event ID.
	fn sse_with_resume<S, H>(
		&self,
		endpoint: &str,
		body: S,
		options: SseOptions<H>,
		last_event_id: Option<&str>,
	) -> impl Send + Future<Output = Result<EventStream<H::Event>>>
	where
		S: Send + Serialize,
		H: 'static + EventHandler;
}

/// Trait implemented by user code to transform raw SSE frames into domain events.
pub trait EventHandler
where
	Self: Send,
{
	/// Output event type generated by this handler.
	type Event;

	/// Called when an `event:` line is seen.
	fn handle_event(&self, #[allow(unused)] event: &str) -> Result<()> {
		Ok(())
	}

	/// Called when a full `data:` block representing one logical event is ready.
	fn handle_data(&self, data: String) -> Result<Self::Event>;

	/// Called when unexpected non-SSE content is encountered.
	fn handle_unexpected(&self, #[allow(unused)] unexpected: String) -> Result<()> {
		Ok(())
	}
}
impl EventHandler for () {
	type Event = String;

	/// Pass-through handler that returns the raw data string.
	fn handle_data(&self, data: String) -> Result<Self::Event> {
		Ok(data)
	}
}

/// Configuration struct that bundles options for an SSE connection.
#[derive(Debug)]
pub struct SseOptions<H> {
	/// If true, `event:` lines are ignored and only `data:` is processed.
	pub drop_event: bool,
	/// User-supplied handler that converts raw SSE frames into events.
	pub event_handler: H,
	/// Policy that governs automatic reconnection behaviour.
	pub reconnect: Reconnect,
}
impl<H> SseOptions<H> {
	/// Creates a new `SseOptions` instance with default settings.
	pub fn new(event_handler: H) -> Self {
		Self { drop_event: false, event_handler, reconnect: Reconnect::default() }
	}

	/// Enables or disables dropping of `event:` frames.
	pub fn drop_event(mut self, drop: bool) -> Self {
		self.drop_event = drop;

		self
	}

	/// Replaces the current event handler with `event_handler`.
	pub fn event_handler(mut self, event_handler: H) -> Self {
		self.event_handler = event_handler;

		self
	}

	/// Replaces the reconnection policy with `reconnect`.
	pub fn reconnect(mut self, reconnect: Reconnect) -> Self {
		self.reconnect = reconnect;

		self
	}
}

/// Policy defining how the client should attempt to reconnect to an SSE stream.
#[derive(Debug)]
pub struct Reconnect {
	/// Whether reconnection attempts are performed.
	pub support: bool,
	/// Maximum number of reconnection attempts before giving up.
	pub max_retries: usize,
	/// Delay between reconnection attempts.
	pub retry_interval: Duration,
}
impl Default for Reconnect {
	fn default() -> Self {
		Self { support: false, max_retries: 3, retry_interval: Duration::from_millis(200) }
	}
}

#[derive(Clone, Debug)]
/// Concrete API client that talks to the remote service using `reqwest`.
pub struct Api {
	http: Client,
	auth: Auth,
}
impl Api {
	/// Constructs a new [`Api`] client with the supplied `auth` settings.
	pub fn new(auth: Auth) -> Self {
		let http =
			Client::builder().user_agent("openagent").build().expect("build must succeed; qed");

		Self { http, auth }
	}
}
impl ApiBase for Api {
	fn base_uri(&self) -> &str {
		&self.auth.uri
	}

	async fn get(&self, endpoint: &str) -> Result<String> {
		Ok(self
			.http
			.get(format!("{}{endpoint}", self.base_uri()))
			.bearer_auth(&self.auth.key)
			.send()
			.await?
			.text()
			.await?)
	}

	async fn post_multipart(&self, endpoint: &str, multipart: Multipart) -> Result<String> {
		Ok(self
			.http
			.post(format!("{}{endpoint}", self.base_uri()))
			.bearer_auth(&self.auth.key)
			.multipart(multipart.into())
			.send()
			.await?
			.text()
			.await?)
	}

	async fn post_json<S>(&self, endpoint: &str, body: S) -> Result<String>
	where
		S: Send + Serialize,
	{
		Ok(self
			.http
			.post(format!("{}{endpoint}", self.base_uri()))
			.bearer_auth(&self.auth.key)
			.json(&body)
			.send()
			.await?
			.text()
			.await?)
	}

	async fn sse<S, H>(
		&self,
		endpoint: &str,
		body: S,
		options: SseOptions<H>,
	) -> Result<Pin<Box<dyn Send + Stream<Item = Result<H::Event>>>>>
	where
		S: Send + Serialize,
		H: 'static + EventHandler,
	{
		let stream = self
			.http
			.post(format!("{}{endpoint}", self.base_uri()))
			.bearer_auth(&self.auth.key)
			.header("Accept", "text/event-stream")
			.header("Cache-Control", "no-cache")
			.json(&body)
			.send()
			.await?
			.bytes_stream()
			.map_err(IoError::other);
		let reader = StreamReader::new(Box::pin(stream) as _);
		let stream = FramedRead::new(reader, LinesCodec::new());

		Ok(Box::pin(Sse {
			stream,
			options,
			last_event: Default::default(),
			data: Default::default(),
			unexpected: Default::default(),
		}))
	}

	async fn sse_with_resume<S, H>(
		&self,
		endpoint: &str,
		body: S,
		options: SseOptions<H>,
		last_event_id: Option<&str>,
	) -> Result<Pin<Box<dyn Send + Stream<Item = Result<H::Event>>>>>
	where
		S: Send + Serialize,
		H: 'static + EventHandler,
	{
		let mut req = self
			.http
			.post(format!("{}{endpoint}", self.base_uri()))
			.bearer_auth(&self.auth.key)
			.header("Accept", "text/event-stream")
			.header("Cache-Control", "no-cache")
			.json(&body);

		// Add Last-Event-ID header for resumption.
		if let Some(event_id) = last_event_id {
			req = req.header("Last-Event-ID", event_id);
		}

		let stream = req.send().await?.bytes_stream().map_err(IoError::other);
		let reader = StreamReader::new(Box::pin(stream) as _);
		let stream = FramedRead::new(reader, LinesCodec::new());

		Ok(Box::pin(Sse {
			stream,
			options,
			last_event: (None, last_event_id.map(Into::into)),
			data: Default::default(),
			unexpected: Default::default(),
		}))
	}
}

/// Authentication tuple holding the API base URL and bearer token.
#[derive(Clone, Debug)]
pub struct Auth {
	/// Remote service root endpoint URL.
	pub uri: String,
	/// Secret authentication key used as bearer token.
	pub key: String,
}
impl Auth {
	/// Builds an `Auth` from the `OPENAI_BASE_URL` and `OPENAI_API_KEY` env variables.
	pub fn from_env() -> Self {
		Auth {
			uri: env::var("OPENAI_BASE_URL").expect("OPENAI_BASE_URL must be set; qed"),
			key: env::var("OPENAI_API_KEY").expect("OPENAI_API_KEY must be set; qed"),
		}
	}
}

/// Helper struct for building multipart/form-data request bodies.
#[derive(Clone, Debug, Default)]
pub struct Multipart {
	/// List of binary parts each containing a name, data buffer, and optional filename.
	#[allow(clippy::type_complexity)]
	pub binary: Vec<(Cow<'static, str>, Cow<'static, [u8]>, Option<String>)>,
	/// List of text parts each containing a name and UTF-8 string value.
	pub text: Vec<(Cow<'static, str>, Cow<'static, str>)>,
}
impl From<Multipart> for Form {
	fn from(val: Multipart) -> Form {
		val.binary.into_iter().fold(
			val.text.into_iter().fold(Form::new(), |form, (k, v)| form.text(k, v)),
			|form, (k, v, filename)| {
				let len = v.len() as _;

				form.part(
					k,
					match v {
						Cow::Borrowed(v) => build_stream_part(v, len, filename),
						Cow::Owned(v) => build_stream_part(v, len, filename),
					},
				)
			},
		)
	}
}

/// Stream wrapper that parses raw bytes from the HTTP response into SSE frames.
#[pin_project::pin_project]
pub struct Sse<T> {
	/// Line-based parser around the raw HTTP byte stream.
	#[pin]
	pub stream: FramedRead<StreamReader<ByteStream, Bytes>, LinesCodec>,
	/// Configuration options controlling behaviour of the SSE consumer.
	pub options: SseOptions<T>,
	/// Tuple storing the most recently observed `(event_type, event_id)`.
	pub last_event: (Option<String>, Option<String>),
	/// Buffer holding concatenated `data:` lines until an empty line finalises the event.
	pub data: String,
	/// Buffer holding non-SSE content encountered in the stream.
	pub unexpected: String,
}
impl<T> Stream for Sse<T>
where
	T: EventHandler,
{
	type Item = Result<T::Event>;

	/// Polls the underlying byte stream and emits parsed events.
	fn poll_next(self: Pin<&mut Self>, ctx: &mut Context) -> Poll<Option<Self::Item>> {
		let mut this = self.project();

		loop {
			match Pin::new(&mut this.stream).poll_next(ctx) {
				Poll::Ready(Some(Ok(line))) => {
					let line = line.trim();

					// Handle SSE protocol.
					if line.is_empty() {
						// Empty line indicates end of an event.
						if !this.data.is_empty() {
							let data = mem::take(this.data);

							// Shrink capacity to free unused memory if the string was large.
							this.data.shrink_to_fit();

							let res = this.options.event_handler.handle_data(data);

							// Clear current event type.
							this.last_event.0 = None;

							return Poll::Ready(Some(res));
						}

						continue;
					}

					tracing::debug!("{line}");

					// Parse SSE line.
					if let Some(data_chunk) = line.strip_prefix("data: ") {
						if data_chunk == "[DONE]" {
							return Poll::Ready(None);
						}

						// Accumulate data.
						if !this.data.is_empty() {
							this.data.push('\n');
						}

						this.data.push_str(data_chunk);
					} else if let Some(event) = line.strip_prefix("event: ") {
						// Handle event.
						if !this.options.drop_event {
							this.last_event.0 = Some(event.into());

							if let Err(e) = this.options.event_handler.handle_event(event) {
								return Poll::Ready(Some(Err(e)));
							}
						}
					} else if let Some(event_id) = line.strip_prefix("id: ") {
						// Store event ID for reconnection.
						this.last_event.1 = Some(event_id.into());
					} else if let Some(retry_ms) = line.strip_prefix("retry: ") {
						// Handle retry instruction (optional implementation).
						if let Ok(_ms) = retry_ms.parse::<u64>() {
							// Update retry interval if needed (currently ignored).
						}
					} else if line.starts_with(':') {
						// Comment line, ignore.
						continue;
					} else {
						// Non-SSE formatted line - accumulate as unexpected content.
						if !this.unexpected.is_empty() {
							this.unexpected.push('\n');
						}

						this.unexpected.push_str(line);
					}
				},
				Poll::Ready(Some(Err(e))) => return Poll::Ready(Some(Err(e.into()))),
				Poll::Ready(None) => {
					// Stream ended - check if we have accumulated unexpected content to process.
					if !this.unexpected.is_empty() {
						let unexpected = mem::take(this.unexpected);

						this.unexpected.shrink_to_fit();

						if let Err(e) = this.options.event_handler.handle_unexpected(unexpected) {
							return Poll::Ready(Some(Err(e)));
						}
					}

					return Poll::Ready(None);
				},
				Poll::Pending => return Poll::Pending,
			}
		}
	}
}

/// Builds a `Part` from raw bytes and an optional filename for multipart uploads.
fn build_stream_part<T>(data: T, data_len: u64, filename: Option<String>) -> Part
where
	T: Into<Body>,
{
	let part = Part::stream_with_length(data, data_len);

	if let Some(filename) = filename { part.file_name(filename) } else { part }
}