procyon 0.1.2

//! Reassembles a `/chat/completions` stream into content blocks.
//!
//! This dialect streams a single choice as deltas: text arrives in `delta.content`, and a tool
//! call arrives as fragments keyed by index — the index is what ties a later argument fragment to
//! the name that opened it.

use std::collections::BTreeMap;

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

use super::wire::{to_neutral_stop_reason, tool_block, StreamChunk};
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>,
    text: String,
    calls: BTreeMap<usize, PartialCall>,
    finish_reason: Option<String>,
    usage: Option<TokenUsage>,
}

impl<'a> StreamState<'a> {
    pub fn new(update_tx: &'a mpsc::UnboundedSender<String>) -> Self {
        Self {
            update_tx,
            text: String::new(),
            calls: BTreeMap::new(),
            finish_reason: None,
            usage: None,
        }
    }

    pub fn into_outcome(self) -> StreamOutcome {
        let mut blocks = Vec::new();
        if !self.text.is_empty() {
            blocks.push(ContentPart::Text { text: self.text });
        }
        for (_, call) in self.calls {
            blocks.push(call.into_block());
        }

        StreamOutcome {
            blocks,
            stop_reason: self.finish_reason.as_deref().map(to_neutral_stop_reason),
            usage: self.usage,
        }
    }
}

impl EventSink for StreamState<'_> {
    fn absorb(&mut self, payload: &str) -> Result<Flow> {
        if payload == "[DONE]" {
            return Ok(Flow::Continue);
        }

        // A provider may interleave keepalives or 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::<StreamChunk>(payload) else {
            return Ok(Flow::Continue);
        };

        // A failure reported mid-stream is a failed request, not an empty answer — the same
        // treatment the Anthropic adapter gives its `error` event.
        if let Some(error) = event.error {
            bail!("Stream error: {}", error);
        }

        if let Some(wire_usage) = event.usage {
            let reported: TokenUsage = wire_usage.into();
            // Some servers attach a usage object to every chunk and fill it only on the last one.
            // Overwriting unconditionally would replace a real figure with zeros and leave the
            // budget estimator anchored on nothing.
            if reported.total() > 0 {
                self.usage = Some(reported);
            }
        }

        let Some(choice) = event.choices.into_iter().next() else {
            return Ok(Flow::Continue);
        };

        if let Some(reason) = choice.finish_reason {
            self.finish_reason = Some(reason);
        }

        if let Some(content) = choice.delta.content.filter(|c| !c.is_empty()) {
            // The receiver is the UI; it hanging up means nobody is reading the rest.
            if self.update_tx.send(content.clone()).is_err() {
                return Ok(Flow::Stop);
            }
            self.text.push_str(&content);
        }

        for fragment in choice.delta.tool_calls.into_iter().flatten() {
            let call = self.calls.entry(fragment.index).or_default();
            if let Some(id) = fragment.id {
                call.id = id;
            }
            if let Some(function) = fragment.function {
                if let Some(name) = function.name {
                    call.name = name;
                }
                if let Some(arguments) = function.arguments {
                    call.arguments.push_str(&arguments);
                }
            }
        }

        Ok(Flow::Continue)
    }
}

#[derive(Debug, Default)]
struct PartialCall {
    id: String,
    name: String,
    arguments: String,
}

impl PartialCall {
    fn into_block(self) -> ContentPart {
        tool_block(self.id, self.name, &self.arguments)
    }
}

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

    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()
    }

    fn run_events(events: &[&str]) -> (StreamOutcome, Vec<String>) {
        let chunks: Vec<Vec<u8>> = events.iter().map(|e| event(e)).collect();
        let refs: Vec<&[u8]> = chunks.iter().map(|c| c.as_slice()).collect();
        run(&refs).unwrap()
    }

    #[test]
    fn content_deltas_assemble_and_stream_as_they_arrive() {
        let (outcome, streamed) = run_events(&[
            r#"{"choices":[{"delta":{"content":"he"}}]}"#,
            r#"{"choices":[{"delta":{"content":"llo"}}]}"#,
            r#"{"choices":[{"delta":{},"finish_reason":"stop"}]}"#,
            "[DONE]",
        ]);

        assert_eq!(
            outcome.blocks,
            vec![ContentPart::Text {
                text: "hello".to_string()
            }]
        );
        assert_eq!(streamed, vec!["he", "llo"]);
        assert_eq!(outcome.stop_reason.as_deref(), Some("end_turn"));
    }

    // The index is the only thing tying an argument fragment to the name that opened the call,
    // and providers send the name once, in the first fragment.
    #[test]
    fn tool_call_fragments_assemble_by_index() {
        let (outcome, _) = run_events(&[
            r#"{"choices":[{"delta":{"tool_calls":[{"index":0,"id":"c1","function":{"name":"read","arguments":"{\"path\":"}}]}}]}"#,
            r#"{"choices":[{"delta":{"tool_calls":[{"index":0,"function":{"arguments":"\"a.rs\"}"}}]}}]}"#,
            r#"{"choices":[{"delta":{},"finish_reason":"tool_calls"}]}"#,
        ]);

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

    // Parallel calls interleave their fragments, so keying by index is what keeps two calls from
    // merging into one.
    #[test]
    fn interleaved_parallel_calls_stay_separate() {
        let (outcome, _) = run_events(&[
            r#"{"choices":[{"delta":{"tool_calls":[{"index":0,"id":"c1","function":{"name":"read","arguments":"{}"}}]}}]}"#,
            r#"{"choices":[{"delta":{"tool_calls":[{"index":1,"id":"c2","function":{"name":"list","arguments":"{"}}]}}]}"#,
            r#"{"choices":[{"delta":{"tool_calls":[{"index":1,"function":{"arguments":"}"}}]}}]}"#,
        ]);

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

    #[test]
    fn usage_is_read_from_the_final_chunk() {
        let (outcome, _) = run_events(&[
            r#"{"choices":[{"delta":{"content":"hi"}}]}"#,
            r#"{"choices":[],"usage":{"prompt_tokens":1000,"completion_tokens":50,"prompt_cache_hit_tokens":800}}"#,
        ]);

        let usage = outcome.usage.expect("usage reported");
        assert_eq!(usage.input, 200);
        assert_eq!(usage.cache_read, 800);
        assert_eq!(usage.output, 50);
    }

    #[test]
    fn a_stream_without_usage_reports_none() {
        let (outcome, _) = run_events(&[r#"{"choices":[{"delta":{"content":"hi"}}]}"#]);
        assert!(outcome.usage.is_none());
    }

    // Servers that attach a usage object to every chunk fill it only on the last one. Letting the
    // zeros win would leave the budget estimator anchored on nothing.
    #[test]
    fn an_empty_usage_chunk_does_not_erase_a_real_one() {
        let (outcome, _) = run_events(&[
            r#"{"choices":[],"usage":{"prompt_tokens":1000,"completion_tokens":50}}"#,
            r#"{"choices":[],"usage":{"prompt_tokens":0,"completion_tokens":0}}"#,
        ]);

        let usage = outcome.usage.expect("usage reported");
        assert_eq!(usage.input, 1000);
        assert_eq!(usage.output, 50);
    }

    // A local server that fails after sending its 200 header reports it as a chunk. Parsed as a
    // chunk with no choices, it used to be indistinguishable from an empty answer.
    #[test]
    fn an_error_chunk_fails_the_request() {
        let chunks = [event(
            r#"{"error":{"message":"model not loaded","code":500}}"#,
        )];
        let refs: Vec<&[u8]> = chunks.iter().map(|c| c.as_slice()).collect();

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

    // Some servers spell it as a bare string rather than an object.
    #[test]
    fn a_bare_string_error_also_fails_the_request() {
        let chunks = [event(r#"{"error":"context length exceeded"}"#)];
        let refs: Vec<&[u8]> = chunks.iter().map(|c| c.as_slice()).collect();

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

    // Providers that never emit parallel calls omit `index`. The whole chunk used to fail to
    // deserialize, and the guard against unmodelled payloads then dropped the call in silence.
    #[test]
    fn a_tool_call_without_an_index_is_still_assembled() {
        let (outcome, _) = run_events(&[
            r#"{"choices":[{"delta":{"tool_calls":[{"id":"c1","function":{"name":"read","arguments":"{\"path\":\"a.rs\"}"}}]}}]}"#,
        ]);

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

    // Local servers and proxies interleave payloads this adapter does not model; abandoning the
    // response over one of them would lose the whole reply.
    #[test]
    fn an_unparseable_payload_is_skipped() {
        let (outcome, _) = run_events(&[
            "not json at all",
            r#"{"choices":[{"delta":{"content":"ok"}}]}"#,
        ]);

        assert_eq!(
            outcome.blocks,
            vec![ContentPart::Text {
                text: "ok".to_string()
            }]
        );
    }

    #[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#"{"choices":[{"delta":{"content":"hi"}}]}"#))
            .unwrap();

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