feat: stuff

src/lib.rs

@@ -0,0 +1,69 @@
+use clap::Parser;
+use secrecy::SecretString;
+use std::{convert::Infallible, net::IpAddr, str::FromStr};
+use tracing_subscriber::EnvFilter;
+
+pub mod openai;
+pub mod server;
+
+#[derive(Debug, Clone)]
+pub struct CloudflareCredentials {
+    pub account_id: String,
+    pub api_key: SecretString,
+}
+
+impl FromStr for CloudflareCredentials {
+    type Err = Infallible; // TODO
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        let (account_id, api_key) = s.split_once(":").expect("TODO");
+
+        let account_id = account_id.into();
+        let api_key = api_key.into();
+
+        Ok(Self {
+            account_id,
+            api_key,
+        })
+    }
+}
+
+#[derive(Debug, Clone, Parser)]
+pub struct AppArgs {
+    #[arg(long, env, default_value_t = [127, 0, 0, 1].into())]
+    pub ip: IpAddr,
+
+    #[arg(long, env, default_value_t = 3126)]
+    pub port: u16,
+
+    #[arg(long)]
+    pub cloudflare_workers_ai_credentials: Vec<CloudflareCredentials>,
+
+    #[arg(long)]
+    pub google_ai_studio_api_key: Vec<SecretString>,
+
+    #[arg(long)]
+    pub groq_api_key: Vec<SecretString>,
+
+    #[arg(long)]
+    pub hugging_face_api_key: Vec<SecretString>,
+
+    #[arg(long)]
+    pub mistral_api_key: Vec<SecretString>,
+
+    #[arg(long)]
+    pub openrouter_api_key: Vec<SecretString>,
+}
+
+#[derive(Parser)]
+pub struct Args {
+    #[arg(
+        long = "tracing-directives",
+        env = "RUST_LOG",
+        default_value = "warn,katniss_ai=debug"
+    )]
+    pub env_filter: EnvFilter,
+
+    #[command(flatten)]
+    pub app_args: AppArgs,
+}

src/main.rs

@@ -1,55 +1,11 @@
-use std::net::{IpAddr, SocketAddr};
-
 use clap::Parser;
-use futures::StreamExt as _;
+use std::net::SocketAddr;
-use reqwest_sse::EventSource;
+use tracing_subscriber::fmt::format::FmtSpan;
-use secrecy::SecretString;
-use tokio::{join, pin, task::LocalSet};
-use tracing_subscriber::{EnvFilter, fmt::format::FmtSpan};
-
-mod openai;
-mod server;
-
-#[derive(Parser, Debug)]
-struct AppArgs {
-    #[arg(long, env, default_value_t = [127, 0, 0, 1].into())]
-    ip: IpAddr,
-
-    #[arg(long, env, default_value_t = 3126)]
-    port: u16,
-
-    #[arg(long)]
-    cloudflare_workers_ai_api_key: Vec<SecretString>,
-
-    #[arg(long)]
-    google_ai_studio_api_key: Vec<SecretString>,
-
-    #[arg(long)]
-    groq_api_key: Vec<SecretString>,
-
-    #[arg(long)]
-    hugging_face_api_key: Vec<SecretString>,
-
-    #[arg(long)]
-    mistral_api_key: Vec<SecretString>,
-
-    #[arg(long)]
-    openrouter_api_key: Vec<SecretString>,
-}
-
-#[derive(Parser)]
-struct Args {
-    #[arg(long = "tracing-directives", env = "RUST_LOG", default_value = "warn,katniss_ai=debug")]
-    env_filter: EnvFilter,
-
-    #[command(flatten)]
-    app_args: AppArgs,
-}
 
 #[tokio::main]
 async fn main() {
     let args = Parser::parse();
-    let Args {
+    let katniss_ai::Args {
         env_filter,
         app_args,
     } = args;
@@ -61,54 +17,20 @@ async fn main() {
         .init();
     tracing::debug!(?app_args, "using");
 
-    rustls::crypto::aws_lc_rs::default_provider().install_default().unwrap();
+    rustls::crypto::aws_lc_rs::default_provider()
+        .install_default()
+        .unwrap();
 
     let addr = SocketAddr::new(app_args.ip, app_args.port);
 
-    let (tx, rx) = async_channel::unbounded();
+    let app = katniss_ai::server::create_app(
-
+        katniss_ai::server::AppState::builder()
-    let app = server::create_app(
-        server::AppState::builder()
             .app_args(app_args.into())
-            .stream_requester(tx)
             .build(),
     );
 
     let listener = tokio::net::TcpListener::bind(addr).await.expect("TODO");
 
-    let stream_responses = async move {
-        let local_set = LocalSet::new();
-
-        local_set
-            .run_until(async {
-                pin!(rx);
-
-                while let Some((response, sse_sink)) = rx.next().await {
-                    local_set.spawn_local(async move {
-                        match response.events().await {
-                            Ok(mut events) => {
-                                while let Some(event_result) = events.next().await {
-                                    let message = event_result
-                                        .map(|event| {
-                                            axum::response::sse::Event::default().data(event.data)
-                                        })
-                                        .map_err(Into::into);
-
-                                    let _ = sse_sink.send(message).await;
-                                }
-                            }
-                            Err(event_source_error) => {
-                                let _ = sse_sink.send(Err(event_source_error.into())).await;
-                            }
-                        };
-                    });
-                }
-            })
-            .await;
-    };
-
-    join!(stream_responses, async move {
     tracing::info!(?addr, "listening");
     axum::serve(listener, app).await.expect("TODO");
-    });
 }

src/openai/chat/completion.rs

@@ -1,7 +1,7 @@
 use serde::{Deserialize, Serialize};
 use serde_with::{DeserializeFromStr, SerializeDisplay};
-use strum::{EnumString, VariantArray};
 use std::collections::HashMap;
+use strum::{EnumString, VariantArray};
 
 // ============================================================================
 // Request: Create Chat Completion
@@ -918,7 +918,9 @@ pub enum PromptCacheRetention {
 }
 
 /// Constrains reasoning effort for reasoning models.
-#[derive(Clone, Debug, EnumString, DeserializeFromStr, strum::Display, SerializeDisplay, VariantArray)]
+#[derive(
+    Clone, Debug, EnumString, DeserializeFromStr, strum::Display, SerializeDisplay, VariantArray,
+)]
 #[strum(serialize_all = "lowercase")]
 pub enum ReasoningEffort {
     /// Disable reasoning.

src/server.rs

@@ -1,4 +1,5 @@
 use crate::AppArgs;
+use crate::CloudflareCredentials;
 use crate::openai;
 use axum::Json;
 use axum::Router;
@@ -7,9 +8,9 @@ use axum::http::StatusCode;
 use axum::response::IntoResponse;
 use axum::response::Response;
 use axum::response::Sse;
-use axum_extra::routing::RouterExt;
+use axum_extra::routing::{RouterExt, TypedPath};
-use axum_extra::routing::TypedPath;
 use ecow::EcoString;
+use futures::StreamExt;
 use headers::Header;
 use headers::authorization::Bearer;
 use itertools::iproduct;
@@ -17,31 +18,26 @@ use mitsein::vec1::Vec1;
 use rand::Rng;
 use rand::RngExt;
 use reqwest::Certificate;
+use reqwest_sse::EventSource as _;
 use secrecy::ExposeSecret;
 use secrecy::SecretString;
 use serde::Deserialize;
-use serde::Serialize;
 use serde_with::DeserializeFromStr;
 use serde_with::SerializeDisplay;
-use strum::EnumString;
-use strum::VariantArray;
 use std::collections::VecDeque;
 use std::convert::Infallible;
 use std::fmt::Display;
 use std::str::FromStr;
 use std::sync::Arc;
+use strum::EnumString;
+use strum::VariantArray;
+use tokio::runtime::Handle;
 use typed_builder::TypedBuilder;
 use url::Url;
 
-type SSEEventResult = Result<axum::response::sse::Event, Box<dyn std::error::Error + Sync + Send>>;
-type StreamRequest = (reqwest::Response, async_channel::Sender<SSEEventResult>);
-type StreamRequester = async_channel::Sender<StreamRequest>;
-
 #[derive(Debug, Clone, TypedBuilder)]
 pub struct AppState {
     app_args: Arc<AppArgs>,
-
-    stream_requester: StreamRequester,
 }
 
 #[derive(Deserialize, TypedPath)]
@@ -58,7 +54,9 @@ enum Provider {
     OpenRouter,
 }
 
-#[derive(Clone, Debug, EnumString, DeserializeFromStr, strum::Display, SerializeDisplay, VariantArray)]
+#[derive(
+    Clone, Debug, EnumString, DeserializeFromStr, strum::Display, SerializeDisplay, VariantArray,
+)]
 #[strum(serialize_all = "lowercase")]
 enum Shortcut {
     XS,
@@ -70,7 +68,6 @@ enum Shortcut {
     Experimental,
 }
 
-
 #[derive(Debug, Clone, DeserializeFromStr, SerializeDisplay)]
 struct ShortcutAndMaybeReasoning(Shortcut, Option<openai::chat::completion::ReasoningEffort>);
 
@@ -106,7 +103,6 @@ impl Display for ShortcutAndMaybeReasoning {
     }
 }
 
-
 impl Shortcut {
     fn suggestions(&self) -> Vec<(Provider, EcoString)> {
         use Provider as P;
@@ -139,6 +135,7 @@ impl Shortcut {
                 (P::HuggingFace, "Qwen/Qwen3.8-27B".into()),
                 (P::HuggingFace, "meta-models/Muse-Glimmer-30B".into()),
                 (P::OpenRouter, "google/gemma-4-31b-it:free".into()),
+                (P::HuggingFace, "google/gemma-4-31B-it".into()),
                 (P::Mistral, "mistral-small-2603".into()),
                 (P::Groq, "groq/compound".into()),
                 (P::OpenRouter, "openrouter/fusion".into()),
@@ -150,6 +147,10 @@ impl Shortcut {
                 (P::Mistral, "mistral-medium-3.5".into()),
                 (P::OpenRouter, "inclusionai/ling-3.0-flash-vl:free".into()),
                 (P::OpenRouter, "poolside/laguna-s-2.1:free".into()),
+                (
+                    P::OpenRouter,
+                    "inclusionai/ling-3.0-flash-sante:free".into(),
+                ),
                 (
                     P::OpenRouter,
                     "nvidia/nemotron-3-super-120b-a12b:free".into(),
@@ -162,28 +163,29 @@ impl Shortcut {
             S::L => vec![
                 (P::CloudflareWorkersAi, "@cf/zai-org/glm-5.3-flash".into()),
                 (P::HuggingFace, "zai-org/GLM-5.3-Flash".into()),
+                (P::HuggingFace, "zai-org/GLM-5.3-Flash-BF16".into()),
                 (P::HuggingFace, "deepseek-ai/DeepSeek-V4.1-Flash".into()),
-                (P::OpenRouter, "nex-agi/nex-n2.5-pro:free".into()),
+                (P::HuggingFace, "XiaomiMiMo/MiMo-V2.6-Flash-RL".into()),
+                (P::OpenRouter, "inclusionai/ling-3.1-flash".into()),
+                (P::OpenRouter, "dots-studio/dots-3-note-preview:free".into()),
                 (P::HuggingFace, "MiniMaxAI/MiniMax-M3".into()),
                 (P::OpenRouter, "thinkingmachines/inkling-small:free".into()),
                 (P::HuggingFace, "thinkingmachines/Inkling-Small".into()),
             ],
 
             S::XL => vec![
-                (P::CloudflareWorkersAi, "@cf/zai-org/glm-5.3".into()),
                 (P::HuggingFace, "zai-org/GLM-5.3".into()),
+                (P::CloudflareWorkersAi, "@cf/zai-org/glm-5.3".into()),
                 (P::Mistral, "zai-glm-5-3".into()),
+                (P::HuggingFace, "tencent/Hy4-preview".into()),
+                (P::HuggingFace, "XiaomiMiMo/MiMo-V2.6-Pro-RL".into()),
                 (P::HuggingFace, "moonshotai/Kimi-K3".into()),
                 (P::HuggingFace, "Qwen/Qwen3.8-2.4T-A95B".into()),
                 (P::OpenRouter, "thinkingmachines/inkling:free".into()),
                 (P::HuggingFace, "thinkingmachines/Inkling".into()),
-                (P::HuggingFace, "tencent/Hy4-preview".into()),
-                (P::HuggingFace, "deepseek-ai/DeepSeek-V4-Pro".into()),
             ],
 
-            S::Experimental => vec![
+            S::Experimental => vec![(P::OpenRouter, "space-bunny-alpha".into())],
-                (P::OpenRouter, "space-bunny-alpha".into())
-            ]
         }
     }
 
@@ -225,10 +227,7 @@ struct Upstream {
 
 enum ChatCompletionResponse {
     Immediate(openai::chat::completion::ChatCompletion),
-    Stream {
+    Stream(reqwest::Response),
-        response: reqwest::Response,
-        stream_requester: StreamRequester,
-    },
 }
 
 impl IntoResponse for ChatCompletionResponse {
@@ -237,17 +236,33 @@ impl IntoResponse for ChatCompletionResponse {
             ChatCompletionResponse::Immediate(chat_completion) => {
                 Json(chat_completion).into_response()
             }
-            ChatCompletionResponse::Stream {
+            ChatCompletionResponse::Stream(response) => {
-                response,
+                let (sink, stream) = async_channel::bounded(256);
-                stream_requester,
-            } => {
-                let (sse_sink, sse_stream) = async_channel::bounded(256);
 
-                stream_requester
+                let handle = Handle::current();
-                    .try_send((response, sse_sink))
-                    .expect("this is an unbounded channel");
 
-                Sse::new(sse_stream).into_response()
+                std::thread::spawn(move || {
+                    handle.block_on(async move {
+                        match response.events().await {
+                            Ok(mut events) => {
+                                while let Some(event_result) = events.next().await {
+                                    let message = event_result
+                                        .map(|event| {
+                                            axum::response::sse::Event::default().data(event.data)
+                                        })
+                                        .map_err(Box::<dyn std::error::Error + Send + Sync>::from);
+
+                                    let _ = sink.send(message).await;
+                                }
+                            }
+                            Err(event_source_error) => {
+                                let _ = sink.send(Err(event_source_error.into())).await;
+                            }
+                        };
+                    });
+                });
+
+                Sse::new(stream).into_response()
             }
         }
     }
@@ -270,14 +285,37 @@ impl IntoResponse for ChatCompletionError {
 async fn create_chat_completion(
     _: V1ChatCompletionsPath,
     State(state): State<AppState>,
-    Json(request): Json<openai::chat::completion::CreateChatCompletionRequest<ShortcutAndMaybeReasoning>>,
+    Json(request): Json<
+        openai::chat::completion::CreateChatCompletionRequest<ShortcutAndMaybeReasoning>,
+    >,
 ) -> Result<ChatCompletionResponse, ChatCompletionError> {
     let ShortcutAndMaybeReasoning(shortcut, reasoning_in_id) = request.model.clone();
     let mut model_upstreams = VecDeque::new();
 
     for (provider, model) in shortcut.suggestions() {
-        let upstreams = match provider {
+        let upstreams: Vec<Upstream> =
-            Provider::CloudflareWorkersAi => vec![], // TODO
+            match provider {
+                Provider::CloudflareWorkersAi => state
+                    .app_args
+                    .cloudflare_workers_ai_credentials
+                    .iter()
+                    .cloned()
+                    .map(
+                        |CloudflareCredentials {
+                             account_id,
+                             api_key,
+                         }| {
+                            Upstream::builder()
+                        .base_url(
+                            format!("https://api.cloudflare.com/client/v4/accounts/{account_id}/ai")
+                                .parse()
+                                .expect("TODO"),
+                        )
+                        .api_key(api_key)
+                        .build()
+                        },
+                    )
+                    .collect(),
                 Provider::GoogleAiStudio => state
                     .app_args
                     .google_ai_studio_api_key
@@ -397,7 +435,11 @@ async fn create_chat_completion(
             web_search_options,
         } = request.clone();
 
-        let reasoning_effort = Some(reasoning_effort.or(reasoning_in_id.clone()).unwrap_or(shortcut.default_reasoning()));
+        let reasoning_effort = Some(
+            reasoning_effort
+                .or(reasoning_in_id.clone())
+                .unwrap_or(shortcut.default_reasoning()),
+        );
 
         let request = openai::chat::completion::CreateChatCompletionRequest {
             model,
@@ -443,7 +485,7 @@ async fn create_chat_completion(
             .tls_certs_only(
                 webpki_root_certs::TLS_SERVER_ROOT_CERTS
                     .iter()
-                    .filter_map(|cert| Certificate::from_der(&cert).ok()),
+                    .filter_map(|cert| Certificate::from_der(cert).ok()),
             )
             .build()
             .expect("TODO");
@@ -500,10 +542,7 @@ async fn create_chat_completion(
         };
 
         if request.stream.unwrap_or(false) {
-            return Ok(ChatCompletionResponse::Stream {
+            return Ok(ChatCompletionResponse::Stream(upstream_response));
-                response: upstream_response,
-                stream_requester: state.stream_requester,
-            });
         } else {
             let response = match upstream_response.json().await {
                 Ok(response) => response,
@@ -525,18 +564,24 @@ async fn create_chat_completion(
 struct V1ModelsPath;
 
 #[tracing::instrument]
-async fn list_models(_: V1ModelsPath) -> Json<openai::models::ListModelsResponse<ShortcutAndMaybeReasoning>> {
+async fn list_models(
-    let shortcuts = Shortcut::VARIANTS.into_iter().cloned();
+    _: V1ModelsPath,
-    let reasoning_efforts = openai::chat::completion::ReasoningEffort::VARIANTS.into_iter().cloned().map(Some).chain(None);
+) -> Json<openai::models::ListModelsResponse<ShortcutAndMaybeReasoning>> {
+    let shortcuts = Shortcut::VARIANTS.iter().cloned();
+    let reasoning_efforts = openai::chat::completion::ReasoningEffort::VARIANTS
+        .iter()
+        .cloned()
+        .map(Some)
+        .chain(None);
 
-    let models = iproduct!(shortcuts, reasoning_efforts)
+    let models = iproduct!(shortcuts, reasoning_efforts).map(|(shortcut, reasoning)| {
-        .into_iter()
+        openai::models::Model {
-        .map(|(shortcut, reasoning)| openai::models::Model {
             id: ShortcutAndMaybeReasoning(shortcut, reasoning),
             created: 0,
             object: openai::models::ModelObject::Model,
             owned_by: "nobody".into(),
             shutdown_date: None,
+        }
     });
 
     Json(openai::models::ListModelsResponse {