1use super::http::{error, read_http, respond};
4use super::{ApiKind, Job, MODEL, failure, registry, sessions};
5use crate::control::State;
6use anyhow::Result;
7use serde_json::{Value, json};
8use std::{
9 net::TcpStream,
10 sync::{Arc, mpsc},
11 time::Duration,
12};
13
14pub(super) fn connection(
15 mut stream: TcpStream,
16 queue: &mpsc::SyncSender<Job>,
17 state: &State,
18 requests: &Arc<registry::Registry>,
19 sessions: &Arc<std::sync::Mutex<sessions::Store>>,
20) -> Result<()> {
21 stream.set_read_timeout(Some(Duration::from_secs(15)))?;
22 stream.set_write_timeout(Some(Duration::from_secs(30)))?;
23 stream.set_nodelay(true)?;
24
25 let settings = state.config();
26 let (method, path, body) = match read_http(&mut stream, settings.config.limits.request_bytes) {
27 Ok(r) => r,
28 Err(e) => return error(&mut stream, 400, &e.to_string()),
29 };
30
31 if let Some(result) =
32 session_endpoint(&method, &path, &body, &settings.config, requests, sessions)
33 {
34 return match result {
35 Ok((status, body)) => respond(&mut stream, status, &body),
36 Err(e) => request_error(&mut stream, state, &e),
37 };
38 }
39
40 match (method.as_str(), path.as_str()) {
41 ("GET", "/health") => respond(&mut stream, 200, &json!({"status":"ok"})),
42 ("GET", "/v1/models") => respond(
43 &mut stream,
44 200,
45 &json!({
46 "object":"list", "data":[{"id":MODEL,"object":"model","created":0,"owned_by":"local"}]
47 }),
48 ),
49 ("POST", "/v1/chat/completions" | "/v1/completions") => {
50 let body: Value = match serde_json::from_slice(&body) {
51 Ok(v) => v,
52 Err(e) => return error(&mut stream, 400, &format!("invalid JSON: {e}")),
53 };
54
55 if !body.is_object() {
56 return error(&mut stream, 400, "request must be an object");
57 }
58
59 if body.get("request_id").is_some_and(|v| !v.is_string()) {
60 return error(&mut stream, 400, "request_id must be a string");
61 }
62
63 let ticket = match requests.register(
64 body["request_id"].as_str(),
65 settings.config.limits.queued_requests + settings.config.limits.active_requests,
66 ) {
67 Ok(ticket) => ticket,
68 Err(e) => return request_error(&mut stream, state, &e),
69 };
70
71 if path == "/v1/completions" && !body["session_id"].is_null() {
72 return error(
73 &mut stream,
74 400,
75 "retained sessions require chat completions",
76 );
77 }
78
79 let session = match sessions::Turn::begin(
80 sessions,
81 &body,
82 &ticket.id,
83 settings.config.limits.request_bytes,
84 ) {
85 Ok(turn) => turn,
86 Err(e) => return request_error(&mut stream, state, &e),
87 };
88 let job = Job {
89 stream,
90 body,
91 settings,
92 ticket,
93 session,
94 kind: if path == "/v1/chat/completions" {
95 ApiKind::Chat
96 } else {
97 ApiKind::Completion
98 },
99 };
100
101 state.update(|s| s.queued_requests += 1);
102
103 let (mut job, message) = match queue.try_send(job) {
104 Ok(()) => return Ok(()),
105 Err(mpsc::TrySendError::Full(job)) => (job, "generation queue is full"),
106 Err(mpsc::TrySendError::Disconnected(job)) => (job, "model unavailable"),
107 };
108
109 state.update(|s| {
110 s.queued_requests -= 1;
111 s.rejected_requests += 1;
112 });
113
114 error(&mut job.stream, 503, message)
115 }
116 _ => error(&mut stream, 404, "unknown endpoint"),
117 }
118}
119
120pub(super) fn session_endpoint(
121 method: &str,
122 path: &str,
123 body: &[u8],
124 config: &crate::config::Config,
125 requests: ®istry::Registry,
126 sessions: &std::sync::Mutex<sessions::Store>,
127) -> Option<Result<(u16, Value)>> {
128 if method == "GET" && path == "/v1/requests" {
129 return Some(Ok((200, requests.list())));
130 }
131
132 if let Some(id) = path
133 .strip_prefix("/v1/requests/")
134 .and_then(|p| p.strip_suffix("/cancel"))
135 {
136 if method != "POST" {
137 return None;
138 }
139
140 return Some(Ok((
141 if requests.cancel(id) { 200 } else { 404 },
142 json!({"id":id}),
143 )));
144 }
145
146 if method == "POST" && path == "/v1/sessions" {
147 return Some(
148 serde_json::from_slice(body)
149 .map_err(Into::into)
150 .and_then(|v| {
151 let mut store = sessions.lock().unwrap();
152 let id = store.create(&v, config)?;
153
154 store.show(&id)
155 })
156 .map(|v| (200, v)),
157 );
158 }
159
160 let id = path.strip_prefix("/v1/sessions/")?;
161
162 match method {
163 "GET" => Some(sessions.lock().unwrap().show(id).map(|v| (200, v))),
164 "DELETE" => Some(
165 sessions
166 .lock()
167 .unwrap()
168 .delete(id)
169 .map(|()| (200, json!({"id":id,"deleted":true}))),
170 ),
171 _ => None,
172 }
173}
174
175fn request_error(out: &mut TcpStream, state: &State, cause: &anyhow::Error) -> Result<()> {
176 let status = failure::status(cause);
177
178 state.update(|s| s.rejected_requests += u64::from(status == 503));
179
180 error(out, status, &cause.to_string())
181}