procyon 0.0.1

//! Reassembles a `/v1/messages` stream into content blocks.
//!
//! Anthropic streams a message as indexed blocks: a `content_block_start` opens one, deltas
//! extend it, and `content_block_stop` closes it. Block reassembly is positional here — one block
//! is open at a time — which is what the wire order guarantees today.

use color_eyre::{eyre::bail, Result};
use tokio::sync::mpsc;

use super::wire::{Delta, StreamEvent};
use crate::agent::{ContentPart, StreamOutcome, TokenUsage};
use crate::sse::{EventSink, Flow};

/// What the adapter has assembled so far from the events seen.
pub struct StreamState<'a> {
    update_tx: &'a mpsc::UnboundedSender<String>,
    blocks: Vec<ContentPart>,
    text: String,
    tool_id: String,
    tool_name: String,
    tool_json: String,
    stop_reason: Option<String>,
    usage: TokenUsage,
    saw_usage: bool,
}

impl<'a> StreamState<'a> {
    pub fn new(update_tx: &'a mpsc::UnboundedSender<String>) -> Self {
        Self {
            update_tx,
            blocks: Vec::new(),
            text: String::new(),
            tool_id: String::new(),
            tool_name: String::new(),
            tool_json: String::new(),
            stop_reason: None,
            usage: TokenUsage::default(),
            saw_usage: false,
        }
    }

    pub fn into_outcome(self) -> StreamOutcome {
        StreamOutcome {
            blocks: self.blocks,
            stop_reason: self.stop_reason,
            usage: self.saw_usage.then_some(self.usage),
        }
    }

    /// Closes the open block, whichever kind it is.
    fn close_block(&mut self) {
        if !self.tool_name.is_empty() {
            let input =
                crate::agent::tool_input(&self.tool_name, &std::mem::take(&mut self.tool_json));

            self.blocks.push(ContentPart::ToolUse {
                id: std::mem::take(&mut self.tool_id),
                name: std::mem::take(&mut self.tool_name),
                input,
            });
        } else if !self.text.is_empty() {
            self.blocks.push(ContentPart::Text {
                text: std::mem::take(&mut self.text),
            });
        }
    }
}

impl EventSink for StreamState<'_> {
    fn absorb(&mut self, payload: &str) -> Result<Flow> {
        // A provider or proxy may interleave fields this adapter does not model; a payload that
        // fails to parse is not a reason to abandon the response.
        let Ok(event) = serde_json::from_str::<StreamEvent>(payload) else {
            return Ok(Flow::Continue);
        };

        match event {
            StreamEvent::ContentBlockStart {
                content_block: ContentPart::ToolUse { id, name, .. },
                ..
            } => {
                self.tool_id = id;
                self.tool_name = name;
                self.tool_json.clear();
            }
            StreamEvent::ContentBlockDelta { delta, .. } => match delta {
                Delta::TextDelta { text } => {
                    self.text.push_str(&text);
                    // The receiver is the UI; it hanging up means nobody is reading the rest.
                    if self.update_tx.send(text).is_err() {
                        return Ok(Flow::Stop);
                    }
                }
                Delta::InputJsonDelta { partial_json } => {
                    self.tool_json.push_str(&partial_json);
                }
            },
            StreamEvent::ContentBlockStop { .. } => self.close_block(),
            StreamEvent::MessageStart { message } => {
                if let Some(reported) = message.usage {
                    self.usage.input = reported.input_tokens.unwrap_or(0);
                    self.usage.cache_read = reported.cache_read_input_tokens.unwrap_or(0);
                    self.usage.cache_write = reported.cache_creation_input_tokens.unwrap_or(0);
                    self.usage.output = reported.output_tokens.unwrap_or(0);
                    self.saw_usage = true;
                }
            }
            StreamEvent::MessageDelta {
                delta,
                usage: reported,
            } => {
                self.stop_reason = delta.stop_reason;
                if let Some(reported) = reported {
                    // Only output grows here; the input side is fixed at message_start.
                    if let Some(output) = reported.output_tokens {
                        self.usage.output = output;
                        self.saw_usage = true;
                    }
                }
            }
            StreamEvent::Error { error } => bail!("Stream error: {}", error.message),
            _ => {}
        }

        Ok(Flow::Continue)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::sse::EventReader;

    /// Drives a whole stream through the reader, as the socket would.
    fn run(chunks: &[&[u8]]) -> Result<(StreamOutcome, Vec<String>)> {
        let (tx, mut rx) = mpsc::unbounded_channel();
        let outcome = {
            let mut reader = EventReader::new(StreamState::new(&tx));
            for chunk in chunks {
                if reader.feed(chunk)? == Flow::Stop {
                    break;
                }
            }
            reader.finish()?;
            reader.into_sink().into_outcome()
        };

        drop(tx);
        let mut streamed = Vec::new();
        while let Ok(chunk) = rx.try_recv() {
            streamed.push(chunk);
        }
        Ok((outcome, streamed))
    }

    fn event(json: &str) -> Vec<u8> {
        format!("data: {}\n\n", json).into_bytes()
    }

    #[test]
    fn text_deltas_assemble_into_one_block_and_stream_as_they_arrive() {
        let chunks = [
            event(
                r#"{"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}"#,
            ),
            event(
                r#"{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"he"}}"#,
            ),
            event(
                r#"{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"llo"}}"#,
            ),
            event(r#"{"type":"content_block_stop","index":0}"#),
        ];
        let refs: Vec<&[u8]> = chunks.iter().map(|c| c.as_slice()).collect();

        let (outcome, streamed) = run(&refs).unwrap();

        assert_eq!(
            outcome.blocks,
            vec![ContentPart::Text {
                text: "hello".to_string()
            }]
        );
        assert_eq!(streamed, vec!["he", "llo"], "deltas must reach the UI live");
    }

    #[test]
    fn a_tool_call_assembles_from_its_argument_fragments() {
        let chunks = [
            event(
                r#"{"type":"content_block_start","index":0,"content_block":{"type":"tool_use","id":"t1","name":"read","input":{}}}"#,
            ),
            event(
                r#"{"type":"content_block_delta","index":0,"delta":{"type":"input_json_delta","partial_json":"{\"path\":"}}"#,
            ),
            event(
                r#"{"type":"content_block_delta","index":0,"delta":{"type":"input_json_delta","partial_json":"\"a.rs\"}"}}"#,
            ),
            event(r#"{"type":"content_block_stop","index":0}"#),
        ];
        let refs: Vec<&[u8]> = chunks.iter().map(|c| c.as_slice()).collect();

        let (outcome, _) = run(&refs).unwrap();

        assert_eq!(
            outcome.blocks,
            vec![ContentPart::ToolUse {
                id: "t1".to_string(),
                name: "read".to_string(),
                input: serde_json::json!({"path": "a.rs"}),
            }]
        );
    }

    // A call with no arguments sends no fragments at all; treating that as malformed would fail
    // every zero-argument tool.
    #[test]
    fn a_tool_call_with_no_arguments_gets_an_empty_object() {
        let chunks = [
            event(
                r#"{"type":"content_block_start","index":0,"content_block":{"type":"tool_use","id":"t1","name":"list","input":{}}}"#,
            ),
            event(r#"{"type":"content_block_stop","index":0}"#),
        ];
        let refs: Vec<&[u8]> = chunks.iter().map(|c| c.as_slice()).collect();

        let (outcome, _) = run(&refs).unwrap();

        assert_eq!(
            outcome.blocks,
            vec![ContentPart::ToolUse {
                id: "t1".to_string(),
                name: "list".to_string(),
                input: serde_json::json!({}),
            }]
        );
    }

    // Unparseable arguments must still produce a `tool_use`. Emitting a `tool_result` instead put
    // it inside an assistant turn, which the API rejects on the next request — and since the turn
    // then held no `tool_use` at all, the loop stopped and the model was never told anything. The
    // call goes out with empty arguments and the tool reports the failure through its own result.
    #[test]
    fn unparseable_arguments_still_yield_a_tool_use() {
        let _guard = crate::diag::test_lock();
        crate::diag::drain();

        let chunks = [
            event(
                r#"{"type":"content_block_start","index":0,"content_block":{"type":"tool_use","id":"t1","name":"read","input":{}}}"#,
            ),
            event(
                r#"{"type":"content_block_delta","index":0,"delta":{"type":"input_json_delta","partial_json":"{invalid"}}"#,
            ),
            event(r#"{"type":"content_block_stop","index":0}"#),
        ];
        let refs: Vec<&[u8]> = chunks.iter().map(|c| c.as_slice()).collect();

        let (outcome, _) = run(&refs).unwrap();

        assert_eq!(
            outcome.blocks,
            vec![ContentPart::ToolUse {
                id: "t1".to_string(),
                name: "read".to_string(),
                input: serde_json::json!({}),
            }]
        );

        // The parse error is the only record that anything went wrong, so it must not be silent.
        let warnings = crate::diag::drain();
        assert!(
            warnings.iter().any(|w| w.contains("did not parse")),
            "expected a recorded warning, got {:?}",
            warnings
        );
    }

    #[test]
    fn usage_is_read_from_message_start_and_updated_by_message_delta() {
        let chunks = [
            event(
                r#"{"type":"message_start","message":{"id":"m","usage":{"input_tokens":1200,"cache_read_input_tokens":400,"cache_creation_input_tokens":30,"output_tokens":1}}}"#,
            ),
            event(
                r#"{"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"output_tokens":915}}"#,
            ),
        ];
        let refs: Vec<&[u8]> = chunks.iter().map(|c| c.as_slice()).collect();

        let (outcome, _) = run(&refs).unwrap();
        let usage = outcome.usage.expect("usage reported");

        assert_eq!(usage.input, 1200);
        assert_eq!(usage.cache_read, 400);
        assert_eq!(usage.cache_write, 30);
        assert_eq!(
            usage.output, 915,
            "message_delta must win over message_start"
        );
        assert_eq!(outcome.stop_reason.as_deref(), Some("end_turn"));
    }

    // Reporting no usage at all has to stay distinguishable from reporting zeros: the budget
    // estimator only corrects itself when it has a real anchor.
    #[test]
    fn a_stream_without_usage_reports_none() {
        let chunks = [event(r#"{"type":"message_stop"}"#)];
        let refs: Vec<&[u8]> = chunks.iter().map(|c| c.as_slice()).collect();

        assert!(run(&refs).unwrap().0.usage.is_none());
    }

    #[test]
    fn an_error_event_fails_the_request() {
        let chunks = [event(
            r#"{"type":"error","error":{"message":"overloaded"}}"#,
        )];
        let refs: Vec<&[u8]> = chunks.iter().map(|c| c.as_slice()).collect();

        let err = run(&refs).unwrap_err().to_string();
        assert!(err.contains("overloaded"), "{}", err);
    }

    #[test]
    fn a_payload_this_adapter_does_not_model_is_skipped() {
        let chunks = [
            event(r#"{"type":"ping"}"#),
            event(r#"{"type":"something_new","weird":true}"#),
            event(
                r#"{"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}"#,
            ),
            event(
                r#"{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"ok"}}"#,
            ),
            event(r#"{"type":"content_block_stop","index":0}"#),
        ];
        let refs: Vec<&[u8]> = chunks.iter().map(|c| c.as_slice()).collect();

        let (outcome, _) = run(&refs).unwrap();
        assert_eq!(
            outcome.blocks,
            vec![ContentPart::Text {
                text: "ok".to_string()
            }]
        );
    }

    // The UI hanging up mid-stream is the reader's cue to stop; anything after it is unread.
    #[test]
    fn a_closed_update_channel_stops_the_stream() {
        let (tx, rx) = mpsc::unbounded_channel();
        drop(rx);

        let mut reader = EventReader::new(StreamState::new(&tx));
        let flow = reader
            .feed(&event(
                r#"{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"hi"}}"#,
            ))
            .unwrap();

        assert_eq!(flow, Flow::Stop);
    }
}