use super::framing;
use super::notifications::{self, Batch};
use super::subscriptions::{Service, Subscriptions};
use serde_json::Value;
use std::io::{self, Write};
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex};
struct State<O> {
output: O,
error: Option<String>,
}
pub(super) struct Writer<O> {
state: Mutex<State<O>>,
limit: usize,
warning_only: Arc<AtomicBool>,
subscriptions: Arc<Subscriptions>,
listens: Arc<super::listen::Listens>,
}
impl<O: Write> Writer<O> {
pub const fn new(
output: O,
limit: usize,
warning_only: Arc<AtomicBool>,
subscriptions: Arc<Subscriptions>,
listens: Arc<super::listen::Listens>,
) -> Self {
Self {
state: Mutex::new(State {
output,
error: None,
}),
limit,
warning_only,
subscriptions,
listens,
}
}
pub fn channel(&self, message: Value, subject: &super::listen::Subject) -> io::Result<()> {
let mut state = self
.state
.lock()
.map_err(|_| io::Error::other("the host writer lock is poisoned"))?;
if let Some(words) = &state.error {
return Err(io::Error::other(words.clone()));
}
if !subject.active.load(Ordering::SeqCst) {
return Ok(());
}
let result = self.put(&mut state, message);
if let Err(error) = &result {
state.error = Some(error.to_string());
}
drop(state);
self.finish(result)
}
pub fn check(&self) -> io::Result<()> {
let state = self
.state
.lock()
.map_err(|_| io::Error::other("the host writer lock is poisoned"))?;
state
.error
.as_ref()
.map_or(Ok(()), |words| Err(io::Error::other(words.clone())))
}
pub fn write(
&self,
message: Value,
service: Option<&Service>,
terminal: bool,
) -> io::Result<()> {
let mut state = self
.state
.lock()
.map_err(|_| io::Error::other("the host writer lock is poisoned"))?;
if let Some(words) = &state.error {
return Err(io::Error::other(words.clone()));
}
if let Some(service) = service {
if !terminal && !service.active.load(Ordering::SeqCst) {
return Ok(());
}
if self.warning_only.load(Ordering::SeqCst) && message["params"]["level"] == "info" {
return Ok(());
}
}
let result = self.put(&mut state, message);
if let Err(error) = &result {
state.error = Some(error.to_string());
}
drop(state);
self.finish(result)
}
pub fn record(&self, service: &Service, subject: &str, batch: &Batch) -> io::Result<()> {
let mut state = self
.state
.lock()
.map_err(|_| io::Error::other("the host writer lock is poisoned"))?;
if let Some(words) = &state.error {
return Err(io::Error::other(words.clone()));
}
if !service.active.load(Ordering::SeqCst) {
return Ok(());
}
let result: io::Result<()> = (|| {
if !self.warning_only.load(Ordering::SeqCst) {
for message in &batch.messages {
self.put(&mut state, notifications::message(service, message))?;
}
}
if batch.dropped > 0 {
self.put(
&mut state,
notifications::warning(
service,
subject,
"dropped",
&format!("{} messages on {subject} were not kept", batch.dropped),
),
)?;
}
Ok(())
})();
if let Err(error) = &result {
state.error = Some(error.to_string());
}
drop(state);
self.finish(result)
}
fn put(&self, state: &mut State<O>, mut message: Value) -> io::Result<()> {
fit(&mut message, self.limit)?;
framing::write_message(&mut state.output, &message)
}
fn finish(&self, result: io::Result<()>) -> io::Result<()> {
if let Err(error) = result {
let mut failures = vec![error.to_string()];
if let Err(words) = self.subscriptions.close_all() {
failures.push(words);
}
if let Err(words) = self.listens.close_all() {
failures.push(words);
}
return Err(io::Error::other(failures.join("; ")));
}
Ok(())
}
}
fn size(message: &Value) -> io::Result<usize> {
Ok(serde_json::to_vec(message)?.len().saturating_add(1))
}
fn fit(message: &mut Value, limit: usize) -> io::Result<()> {
if size(message)? <= limit {
return Ok(());
}
let channel = message["method"] == "notifications/claude/channel";
let pointer = if channel {
"/params/content"
} else {
"/params/data/words"
};
let Some(words) = message
.pointer(pointer)
.and_then(Value::as_str)
.map(str::to_owned)
else {
return Err(io::Error::other(format!(
"the host answer exceeds frame_bytes {limit}; its fields were not cut"
)));
};
if channel {
message["params"]["meta"]["truncated"] = Value::String("true".to_owned());
}
*message
.pointer_mut(pointer)
.ok_or_else(|| io::Error::other("the notification text field disappeared"))? =
Value::String("…".to_owned());
if size(message)? > limit {
return Err(io::Error::other(format!(
"the notification fields exceed frame_bytes {limit}; its fields were not cut"
)));
}
let boundaries: Vec<_> = words
.char_indices()
.map(|(index, _)| index)
.chain(std::iter::once(words.len()))
.collect();
let mut low = 0;
let mut high = boundaries.len() - 1;
while low < high {
let mid = low + (high - low).div_ceil(2);
*message
.pointer_mut(pointer)
.ok_or_else(|| io::Error::other("the notification text field disappeared"))? =
Value::String(format!("{}…", &words[..boundaries[mid]]));
if size(message)? <= limit {
low = mid;
} else {
high = mid - 1;
}
}
*message
.pointer_mut(pointer)
.ok_or_else(|| io::Error::other("the notification text field disappeared"))? =
Value::String(format!("{}…", &words[..boundaries[low]]));
Ok(())
}