Complete local HTTP endpoint parity

src/server/response.rs

@@ -7,7 +7,11 @@ pub(super) fn final_response(
     active: crate::runtime::ActiveGeneration,
     id: &str,
 ) -> Result<(), (u16, String)> {
-    let output = wait_for_output(active)?;
+    let output = match wait_for_output(stream, active) {
+        Ok(output) => output,
+        Err(error) if error == "client disconnected" => return Ok(()),
+        Err(error) => return send_generation_error(stream, request.protocol, request.cors, error),
+    };
     let (content, calls) = parse_generated_tools(state, &output.message.content, request.protocol);
     let finish = if calls.is_empty() {
         output.finish_reason
@@ -114,7 +118,7 @@ pub(super) fn final_response(
             })
         }
     };
-    send_json(stream, 200, &body).map_err(|error| (500, error))
+    send_json_with_cors(stream, 200, &body, request.cors).map_err(|error| (500, error))
 }
 
 pub(super) fn stream_response(
@@ -158,7 +162,7 @@ fn completion_stream_response(
     active: crate::runtime::ActiveGeneration,
     id: &str,
 ) -> Result<(), (u16, String)> {
-    send_sse_headers(stream).map_err(|error| (500, error))?;
+    send_sse_headers_with_cors(stream, request.cors).map_err(|error| (500, error))?;
     while let Ok(event) = active.events.recv() {
         match event {
             GenerationEvent::Chunk {
@@ -215,7 +219,7 @@ fn anthropic_stream_response(
     active: crate::runtime::ActiveGeneration,
     id: &str,
 ) -> Result<(), (u16, String)> {
-    send_sse_headers(stream).map_err(|error| (500, error))?;
+    send_sse_headers_with_cors(stream, request.cors).map_err(|error| (500, error))?;
     let mut prompt_tokens = 0;
     let mut started = false;
     let mut block = None::<(usize, bool)>;
@@ -298,7 +302,7 @@ fn anthropic_stream_response(
                     )?;
                 }
                 if request.has_tools {
-                    let events = projector.push("", true, "toolu_");
+                    let events = projector.finish("toolu_");
                     send_anthropic_projection_events(
                         stream,
                         events,
@@ -468,7 +472,7 @@ fn responses_stream_response(
     active: crate::runtime::ActiveGeneration,
     id: &str,
 ) -> Result<(), (u16, String)> {
-    send_sse_headers(stream).map_err(|error| (500, error))?;
+    send_sse_headers_with_cors(stream, request.cors).map_err(|error| (500, error))?;
     let created = unix_time();
     let message_id = random_id("msg_");
     let reasoning_id = random_id("rs_");
@@ -739,7 +743,8 @@ fn stream_response_with_keepalive(
             } => {
                 prefilling = tokens_per_second.is_none();
                 if prefilling && !headers_sent {
-                    send_sse_headers(stream).map_err(|error| (500, error))?;
+                    send_sse_headers_with_cors(stream, request.cors)
+                        .map_err(|error| (500, error))?;
                     headers_sent = true;
                     last_keepalive = Instant::now();
                 } else if !prefilling {
@@ -785,7 +790,12 @@ fn stream_response_with_keepalive(
                             let _ = send_sse_error(stream, &error);
                             return Ok(());
                         }
-                        return Err((500, error));
+                        return send_generation_error(
+                            stream,
+                            request.protocol,
+                            request.cors,
+                            error,
+                        );
                     }
                 }
                 break;
@@ -801,7 +811,7 @@ fn stream_response_with_keepalive(
         None => return Err((500, "The model runtime stopped unexpectedly.".into())),
     };
     if request.has_tools {
-        let events = projector.push("", true, "call_");
+        let events = projector.finish("call_");
         send_chat_projection_events(stream, &request, id, events)?;
     }
     let (_, calls) = parse_generated_tools_with_ids(
@@ -885,7 +895,7 @@ pub(super) fn send_stream_start(
     role_sent: &mut bool,
 ) -> Result<(), (u16, String)> {
     if !*headers_sent {
-        send_sse_headers(stream).map_err(|error| (500, error))?;
+        send_sse_headers_with_cors(stream, request.cors).map_err(|error| (500, error))?;
         *headers_sent = true;
     }
     if !*role_sent {
@@ -924,10 +934,28 @@ pub(super) fn receive_stream_event(
 }
 
 fn wait_for_output(
+    stream: &mut TcpStream,
     active: crate::runtime::ActiveGeneration,
-) -> Result<crate::engine::GenerationOutput, (u16, String)> {
+) -> Result<crate::engine::GenerationOutput, String> {
     let mut content = String::new();
-    while let Ok(event) = active.events.recv() {
+    stream
+        .set_nonblocking(true)
+        .map_err(|error| error.to_string())?;
+    loop {
+        let event = match active.events.recv_timeout(Duration::from_millis(100)) {
+            Ok(event) => event,
+            Err(std::sync::mpsc::RecvTimeoutError::Timeout) => {
+                let mut byte = [0];
+                match stream.peek(&mut byte) {
+                    Ok(0) => return Err("client disconnected".into()),
+                    Ok(_) => {}
+                    Err(error) if error.kind() == std::io::ErrorKind::WouldBlock => {}
+                    Err(error) => return Err(error.to_string()),
+                }
+                continue;
+            }
+            Err(std::sync::mpsc::RecvTimeoutError::Disconnected) => break,
+        };
         match event {
             GenerationEvent::Chunk {
                 reasoning: false,
@@ -942,12 +970,62 @@ fn wait_for_output(
                 }
             }
             GenerationEvent::Finished(result) => {
-                return result.map_err(|error| (500, error));
+                stream
+                    .set_nonblocking(false)
+                    .map_err(|error| error.to_string())?;
+                return result;
             }
             _ => {}
         }
     }
-    Err((500, "The model runtime stopped unexpectedly.".into()))
+    let _ = stream.set_nonblocking(false);
+    Err("The model runtime stopped unexpectedly.".into())
+}
+
+pub(super) fn send_generation_error(
+    stream: &mut TcpStream,
+    protocol: Protocol,
+    cors: bool,
+    error: String,
+) -> Result<(), (u16, String)> {
+    let Some((prompt_tokens, context)) = context_error_dimensions(&error) else {
+        return Err((500, error));
+    };
+    let body = if protocol == Protocol::Anthropic {
+        json!({
+            "type": "error",
+            "error": {
+                "type": "invalid_request_error",
+                "message": error,
+                "n_prompt_tokens": prompt_tokens,
+                "n_ctx": context
+            }
+        })
+    } else {
+        let parameter = match protocol {
+            Protocol::Completion => "prompt",
+            Protocol::Responses => "input",
+            Protocol::Chat => "messages",
+            Protocol::Anthropic => unreachable!(),
+        };
+        json!({"error": {
+            "message": error,
+            "type": "invalid_request_error",
+            "param": parameter,
+            "code": "context_length_exceeded",
+            "n_prompt_tokens": prompt_tokens,
+            "n_ctx": context
+        }})
+    };
+    send_json_with_cors(stream, 400, &body, cors).map_err(|error| (500, error))
+}
+
+pub(super) fn context_error_dimensions(error: &str) -> Option<(u32, u32)> {
+    let values = error
+        .strip_prefix("Prompt has ")?
+        .strip_suffix(" tokens")?
+        .split_once(" tokens, but the configured context size is ")?;
+    Some((values.0.parse().ok()?, values.1.parse().ok()?))
 }
 
 pub(super) fn usage_json(prompt: u32, cached: u32, completion: u32) -> Value {