1use std::collections::{HashMap, HashSet};
32use std::env;
33use std::fs;
34use std::path::Path as StdPath;
35use std::sync::{Arc, RwLock};
36use std::time::{SystemTime, UNIX_EPOCH};
37
38use axum::extract::{Path, State};
39use axum::response::{IntoResponse, Response as AxumResponse};
40use axum::Json;
41use embacle::config::{CliRunnerType, RunnerConfig};
42use embacle::discovery::resolve_binary;
43use embacle::factory::create_runner_with_config;
44use embacle::types::{ChatMessage, ChatRequest, LlmProvider, RunnerError};
45use serde::{Deserialize, Serialize};
46use serde_json::{Map, Value};
47use uuid::fmt::Simple;
48use uuid::Uuid;
49
50use super::error::{ErrorType, UhpFailure};
51use super::files;
52use super::harnesses;
53use super::state::UhpState;
54
55const RESERVED: [&str; 2] = ["tools", "include"];
57
58#[derive(Debug, Clone, Deserialize)]
60pub struct CreateResponse {
61 pub input: Value,
63 #[serde(default)]
65 pub model: Option<String>,
66 #[serde(default)]
68 pub metadata: Option<Map<String, Value>>,
69 #[serde(default)]
71 pub stream: bool,
72 #[serde(default)]
74 pub background: bool,
75 #[serde(default)]
77 pub previous_response_id: Option<String>,
78 #[serde(default)]
80 pub instructions: Option<String>,
81 #[serde(flatten)]
84 pub rest: Map<String, Value>,
85}
86
87impl CreateResponse {
88 pub fn prompt(&self) -> String {
90 match &self.input {
91 Value::String(s) => s.clone(),
92 Value::Array(items) => items
93 .iter()
94 .filter_map(|item| {
95 item.get("content").map(|c| match c {
96 Value::String(s) => s.clone(),
97 other => other
98 .as_array()
99 .map(|parts| {
100 parts
101 .iter()
102 .filter_map(|p| p.get("text").and_then(Value::as_str))
103 .collect::<Vec<_>>()
104 .join("")
105 })
106 .unwrap_or_default(),
107 })
108 })
109 .collect::<Vec<_>>()
110 .join("\n"),
111 other => other.as_str().unwrap_or_default().to_owned(),
112 }
113 }
114
115 pub fn ignored_fields(&self) -> Vec<String> {
121 RESERVED
122 .iter()
123 .filter(|f| self.rest.contains_key(**f))
124 .map(|f| (*f).to_owned())
125 .collect()
126 }
127
128 pub fn harness_id(&self) -> Option<&str> {
130 self.metadata
131 .as_ref()?
132 .get("harness_id")
133 .and_then(Value::as_str)
134 }
135}
136
137#[derive(Debug, Clone, Serialize)]
139pub struct OutputItem {
140 #[serde(rename = "type")]
142 pub kind: &'static str,
143 pub id: String,
145 pub status: &'static str,
147 pub role: &'static str,
149 pub content: Vec<ContentPart>,
151}
152
153#[derive(Debug, Clone, Serialize)]
155pub struct ContentPart {
156 #[serde(rename = "type")]
158 pub kind: &'static str,
159 pub text: String,
161}
162
163#[derive(Debug, Clone, Copy, Serialize)]
165pub struct Usage {
166 pub input_tokens: u32,
168 pub output_tokens: u32,
170 pub total_tokens: u32,
172}
173
174#[derive(Debug, Clone, Serialize)]
176pub struct Response {
177 pub id: String,
179 pub object: &'static str,
181 pub created_at: u64,
183 pub status: &'static str,
185 pub error: Option<super::error::UhpError>,
187 pub previous_response_id: Option<String>,
189 pub model: String,
191 pub output: Vec<OutputItem>,
193 pub store: bool,
195 pub usage: Option<Usage>,
197 pub metadata: Map<String, Value>,
199}
200
201#[derive(Debug, Default)]
206pub struct ResponseStore {
207 responses: RwLock<HashMap<String, Response>>,
208 sessions: RwLock<HashMap<String, Vec<String>>>,
209 cancelled: RwLock<HashSet<String>>,
210}
211
212impl ResponseStore {
213 pub fn put(&self, session_id: &str, response: Response) {
215 if let Ok(mut sessions) = self.sessions.write() {
216 sessions
217 .entry(session_id.to_owned())
218 .or_default()
219 .push(response.id.clone());
220 }
221 if let Ok(mut responses) = self.responses.write() {
222 responses.insert(response.id.clone(), response);
223 }
224 }
225
226 pub fn get(&self, id: &str) -> Option<Response> {
228 self.responses.read().ok()?.get(id).cloned()
229 }
230
231 pub fn session_ids(&self) -> Vec<String> {
233 self.sessions
234 .read()
235 .map(|s| s.keys().cloned().collect())
236 .unwrap_or_default()
237 }
238
239 pub fn session_turns(&self, session_id: &str) -> Option<Vec<Response>> {
241 let ids = self.sessions.read().ok()?.get(session_id)?.clone();
242 let responses = self.responses.read().ok()?;
243 Some(
244 ids.iter()
245 .filter_map(|i| responses.get(i).cloned())
246 .collect(),
247 )
248 }
249
250 pub fn cancel(&self, id: &str) -> bool {
257 let running = self.get(id).is_some_and(|r| r.status == "in_progress");
258 if running {
259 if let Ok(mut cancelled) = self.cancelled.write() {
260 cancelled.insert(id.to_owned());
261 }
262 if let Ok(mut responses) = self.responses.write() {
263 if let Some(r) = responses.get_mut(id) {
264 r.status = "cancelled";
265 }
266 }
267 }
268 running
269 }
270
271 pub fn is_cancelled(&self, id: &str) -> bool {
273 self.cancelled.read().is_ok_and(|c| c.contains(id))
274 }
275
276 pub fn settle(&self, response: Response) {
281 if self.is_cancelled(&response.id) {
282 return;
283 }
284 if let Ok(mut responses) = self.responses.write() {
285 responses.insert(response.id.clone(), response);
286 }
287 }
288
289 pub fn running_in(&self, session_id: &str) -> Vec<String> {
291 self.session_turns(session_id)
292 .unwrap_or_default()
293 .into_iter()
294 .filter(|r| r.status == "in_progress")
295 .map(|r| r.id)
296 .collect()
297 }
298
299 pub fn delete_session(&self, session_id: &str) -> bool {
304 let ids = self
305 .sessions
306 .write()
307 .map_or(None, |mut sessions| sessions.remove(session_id));
308 let Some(ids) = ids else { return false };
309 if let Ok(mut responses) = self.responses.write() {
310 for id in &ids {
311 responses.remove(id);
312 }
313 }
314 true
315 }
316
317 pub fn session_of(&self, response_id: &str) -> Option<String> {
319 self.get(response_id)?
320 .metadata
321 .get("session_id")
322 .and_then(Value::as_str)
323 .map(str::to_owned)
324 }
325}
326
327fn now_secs() -> u64 {
329 SystemTime::now()
330 .duration_since(UNIX_EPOCH)
331 .map(|d| d.as_secs())
332 .unwrap_or_default()
333}
334
335const SESSION_PREFIX: &str = "sess_";
337
338fn id_with(prefix: &str) -> String {
340 format!("{prefix}{}", Uuid::new_v4().simple())
341}
342
343#[must_use]
350pub fn is_session_id(id: &str) -> bool {
351 id.strip_prefix(SESSION_PREFIX).is_some_and(|hex| {
352 hex.len() == Simple::LENGTH
353 && hex
354 .bytes()
355 .all(|b| b.is_ascii_digit() || (b'a'..=b'f').contains(&b))
356 })
357}
358
359pub async fn create(
366 State(state): State<UhpState>,
367 Json(request): Json<CreateResponse>,
368) -> Result<AxumResponse, UhpFailure> {
369 if request.background {
370 return Ok(Json(start_background(state, request)).into_response());
371 }
372 if request.stream {
373 let prepared = prepare(&state, &request).await?;
374 let previous = request.previous_response_id.clone();
375 return Ok(super::streaming::sse(state, prepared, previous).into_response());
376 }
377 let outcome = run(&state, &request).await?;
378 Ok(Json(outcome.response).into_response())
379}
380
381fn start_background(state: UhpState, request: CreateResponse) -> Response {
392 let id = id_with("resp_");
393 let session_id = request
394 .previous_response_id
395 .as_ref()
396 .and_then(|prev| state.store.session_of(prev))
397 .unwrap_or_else(|| id_with(SESSION_PREFIX));
398
399 let mut metadata = request.metadata.clone().unwrap_or_default();
400 metadata.insert("session_id".to_owned(), Value::String(session_id.clone()));
401
402 let pending = Response {
403 id: id.clone(),
404 object: "response",
405 created_at: now_secs(),
406 status: "in_progress",
407 error: None,
408 previous_response_id: request.previous_response_id.clone(),
409 model: request.model.clone().unwrap_or_default(),
410 output: Vec::new(),
411 store: true,
412 usage: None,
413 metadata,
414 };
415 state.store.put(&session_id, pending.clone());
416
417 let spawned = pending.clone();
418 tokio::spawn(async move {
419 let finished = match run(&state, &request).await {
420 Ok(outcome) => {
421 let mut r = outcome.response;
424 r.id = id;
425 r
426 }
427 Err(failure) => {
428 let mut r = spawned;
429 r.status = "failed";
430 r.error = Some(failure.into_error());
431 r
432 }
433 };
434 state.store.settle(finished);
435 });
436
437 pending
438}
439
440pub async fn cancel(
446 State(state): State<UhpState>,
447 Path(id): Path<String>,
448) -> Result<Json<Response>, UhpFailure> {
449 if state.store.get(&id).is_none() {
450 return Err(UhpFailure::not_found(
451 "response_not_found",
452 "no response with that id exists",
453 ));
454 }
455 state.store.cancel(&id);
456 state.store.get(&id).map(Json).ok_or_else(|| {
457 UhpFailure::not_found("response_not_found", "no response with that id exists")
458 })
459}
460
461pub struct Outcome {
463 pub response: Response,
465}
466
467pub struct Prepared {
469 pub runner: Arc<dyn LlmProvider>,
471 pub chat: ChatRequest,
473 pub model: String,
475 pub session_id: String,
477 pub id: String,
479 pub metadata: Map<String, Value>,
481 pub requested_model: Option<String>,
483}
484
485async fn runner_in(
491 provider: CliRunnerType,
492 workdir: &StdPath,
493) -> Result<Arc<dyn LlmProvider>, RunnerError> {
494 fs::create_dir_all(workdir)
495 .map_err(|e| RunnerError::internal(format!("session working folder: {e}")))?;
496
497 let env_override = env::var(provider.env_override_key()).ok();
498 let binary = resolve_binary(provider.binary_name(), env_override.as_deref())?;
499 let mut config = RunnerConfig::new(binary).with_working_directory(workdir.to_path_buf());
500 config.extra_args = write_permission_args(provider);
501
502 Ok(Arc::from(
503 create_runner_with_config(provider, config).await?,
504 ))
505}
506
507fn write_permission_args(provider: CliRunnerType) -> Vec<String> {
528 match provider {
529 CliRunnerType::ClaudeCode => vec![
530 "--permission-mode".to_owned(),
531 "acceptEdits".to_owned(),
532 "--allowed-tools".to_owned(),
533 "Write Edit".to_owned(),
534 ],
535 _ => Vec::new(),
536 }
537}
538
539pub async fn prepare(state: &UhpState, request: &CreateResponse) -> Result<Prepared, UhpFailure> {
548 let harnesses = harnesses::discover(&state.shared).await;
549 let chosen = match request.harness_id() {
550 Some(id) => harnesses
551 .iter()
552 .find(|d| d.harness.id == id)
553 .ok_or_else(|| {
554 UhpFailure::not_found("harness_not_found", "no harness with that id is configured")
555 })?,
556 None => harnesses.first().ok_or_else(|| {
557 UhpFailure::new(
558 ErrorType::InvalidRequestError,
559 "harness_not_found",
560 "this server has no installed harness to run",
561 )
562 })?,
563 };
564
565 let session_id = request
571 .previous_response_id
572 .as_ref()
573 .and_then(|prev| state.store.session_of(prev))
574 .unwrap_or_else(|| id_with(SESSION_PREFIX));
575 let workdir = files::session_workdir(&session_id);
576 let runner = runner_in(chosen.provider, &workdir).await.map_err(|_| {
577 UhpFailure::new(
578 ErrorType::HarnessError,
579 "harness_unavailable",
580 "the harness could not be started",
581 )
582 })?;
583
584 let requested_model = request.model.clone();
585 let model = requested_model
586 .clone()
587 .or_else(|| chosen.harness.default_model.clone())
588 .unwrap_or_else(|| runner.default_model().to_owned());
589
590 let mut chat = ChatRequest::new(vec![ChatMessage::user(request.prompt())]);
591 chat.model = Some(model.clone());
592
593 let mut metadata = request.metadata.clone().unwrap_or_default();
594 metadata.insert("session_id".to_owned(), Value::String(session_id.clone()));
595 let ignored = request.ignored_fields();
596 if !ignored.is_empty() {
597 metadata.insert(
598 "ignored_fields".to_owned(),
599 Value::Array(ignored.into_iter().map(Value::String).collect()),
600 );
601 }
602
603 Ok(Prepared {
604 runner,
605 chat,
606 model,
607 session_id,
608 id: id_with("resp_"),
609 metadata,
610 requested_model,
611 })
612}
613
614#[must_use]
616pub fn pending_response(p: &Prepared, previous: Option<String>) -> Response {
617 Response {
618 id: p.id.clone(),
619 object: "response",
620 created_at: now_secs(),
621 status: "in_progress",
622 error: None,
623 previous_response_id: previous,
624 model: p.model.clone(),
625 output: Vec::new(),
626 store: true,
627 usage: None,
628 metadata: p.metadata.clone(),
629 }
630}
631
632#[must_use]
634pub fn completed_response(
635 p: Prepared,
636 previous: Option<String>,
637 text: String,
638 usage: Option<Usage>,
639 ran_model: String,
640) -> Response {
641 let mut metadata = p.metadata;
642 if p.requested_model.as_ref().is_some_and(|m| *m != ran_model) {
643 metadata.insert(
644 "requested_model".to_owned(),
645 Value::String(p.requested_model.unwrap_or_default()),
646 );
647 metadata.insert("model_fallback".to_owned(), Value::Bool(true));
648 }
649 Response {
650 id: p.id,
651 object: "response",
652 created_at: now_secs(),
653 status: "completed",
654 error: None,
655 previous_response_id: previous,
656 model: ran_model,
657 output: vec![OutputItem {
658 kind: "message",
659 id: id_with("msg_"),
660 status: "completed",
661 role: "assistant",
662 content: vec![ContentPart {
663 kind: "output_text",
664 text,
665 }],
666 }],
667 store: true,
668 usage,
669 metadata,
670 }
671}
672
673pub async fn run(state: &UhpState, request: &CreateResponse) -> Result<Outcome, UhpFailure> {
679 let prepared = prepare(state, request).await?;
680 let session_id = prepared.session_id.clone();
681 let previous = request.previous_response_id.clone();
682 let model = prepared.model.clone();
683
684 let response = match prepared.runner.complete(&prepared.chat).await {
685 Ok(reply) => {
686 let ran = if reply.model.is_empty() {
687 model
688 } else {
689 reply.model
690 };
691 let usage = reply.usage.map(|u| Usage {
694 input_tokens: u.prompt_tokens,
695 output_tokens: u.completion_tokens,
696 total_tokens: u.total_tokens,
697 });
698 completed_response(prepared, previous, reply.content, usage, ran)
699 }
700 Err(e) => {
701 let mut failed = pending_response(&prepared, previous);
702 failed.status = "failed";
703 failed.error = Some(super::error::UhpError {
704 kind: ErrorType::HarnessError,
705 code: "harness_error".to_owned(),
706 message: e.message,
707 detail: None,
708 });
709 failed
710 }
711 };
712
713 state.store.put(&session_id, response.clone());
714 Ok(Outcome { response })
715}
716
717pub async fn get(
723 State(state): State<UhpState>,
724 Path(id): Path<String>,
725) -> Result<Json<Response>, UhpFailure> {
726 state.store.get(&id).map(Json).ok_or_else(|| {
727 UhpFailure::not_found("response_not_found", "no response with that id exists")
728 })
729}
730
731#[cfg(test)]
732mod tests {
733 use super::*;
734 use serde_json::json;
735
736 fn parse(body: Value) -> CreateResponse {
737 serde_json::from_value(body).expect("request parses") }
739
740 #[test]
741 fn a_bare_string_input_is_one_user_message() {
742 assert_eq!(parse(json!({"input": "hello"})).prompt(), "hello");
743 }
744
745 #[test]
746 fn an_item_array_input_is_flattened_to_its_text() {
747 let r = parse(json!({"input": [
748 {"role": "user", "content": [{"type": "input_text", "text": "first"}]},
749 {"role": "user", "content": "second"}
750 ]}));
751 assert_eq!(r.prompt(), "first\nsecond");
752 }
753
754 #[test]
755 fn ignored_fields_reports_only_what_the_request_carried() {
756 let both = parse(json!({"input": "x", "tools": [], "include": []}));
759 assert_eq!(both.ignored_fields(), vec!["tools", "include"]);
760
761 let one = parse(json!({"input": "x", "include": []}));
762 assert_eq!(one.ignored_fields(), vec!["include"]);
763
764 let neither = parse(json!({"input": "x"}));
765 assert!(
766 neither.ignored_fields().is_empty(),
767 "a request that sent no reserved field must not be told one was ignored"
768 );
769 }
770
771 #[test]
772 fn a_reserved_field_does_not_make_the_request_unparseable() {
773 let r = parse(json!({
776 "input": "x",
777 "tools": [{"type": "function", "name": "probe"}],
778 "metadata": {"harness_id": "chrn_copilot"}
779 }));
780 assert_eq!(r.harness_id(), Some("chrn_copilot"));
781 assert_eq!(r.prompt(), "x");
782 }
783
784 #[test]
785 fn ids_carry_the_prefix_the_object_model_requires() {
786 assert!(id_with("resp_").starts_with("resp_"));
787 assert!(id_with(SESSION_PREFIX).starts_with("sess_"));
788 }
789
790 #[test]
791 fn only_the_session_id_shape_this_server_mints_is_accepted() {
792 assert!(is_session_id(&id_with(SESSION_PREFIX)));
793 for refused in [
794 "..",
795 ".",
796 "",
797 "sess_",
798 "/etc",
799 "sess_0123456789abcdef0123456789abcdef/..",
800 "sess_0123456789abcdef0123456789abcde/",
801 "sess_0123456789ABCDEF0123456789ABCDEF",
802 "sess_0123456789abcdef0123456789abcdef0",
803 "resp_0123456789abcdef0123456789abcdef",
804 ] {
805 assert!(!is_session_id(refused), "{refused:?} is not a minted id");
806 }
807 }
808
809 #[test]
810 fn a_terminal_response_reads_back_unchanged() {
811 let store = ResponseStore::default();
812 let resp = Response {
813 id: "resp_abc".to_owned(),
814 object: "response",
815 created_at: 1,
816 status: "completed",
817 error: None,
818 previous_response_id: None,
819 model: "m".to_owned(),
820 output: Vec::new(),
821 store: true,
822 usage: None,
823 metadata: Map::new(),
824 };
825 store.put("sess_1", resp);
826
827 let got = store.get("resp_abc").expect("stored"); assert_eq!(got.status, "completed", "a terminal state must not change");
829 assert_eq!(store.session_turns("sess_1").map(|t| t.len()), Some(1));
830 assert!(store.get("resp_missing").is_none());
831 }
832}