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::invocation::{self, StagedSkills};
54use super::state::UhpState;
55
56const RESERVED: [&str; 2] = ["tools", "include"];
58
59#[derive(Debug, Clone, Deserialize)]
61pub struct CreateResponse {
62 pub input: Value,
64 #[serde(default)]
66 pub model: Option<String>,
67 #[serde(default)]
69 pub metadata: Option<Map<String, Value>>,
70 #[serde(default)]
72 pub stream: bool,
73 #[serde(default)]
75 pub background: bool,
76 #[serde(default)]
78 pub previous_response_id: Option<String>,
79 #[serde(default)]
81 pub instructions: Option<String>,
82 #[serde(flatten)]
85 pub rest: Map<String, Value>,
86}
87
88impl CreateResponse {
89 pub fn prompt(&self) -> String {
91 match &self.input {
92 Value::String(s) => s.clone(),
93 Value::Array(items) => items
94 .iter()
95 .filter_map(|item| {
96 item.get("content").map(|c| match c {
97 Value::String(s) => s.clone(),
98 other => other
99 .as_array()
100 .map(|parts| {
101 parts
102 .iter()
103 .filter_map(|p| p.get("text").and_then(Value::as_str))
104 .collect::<Vec<_>>()
105 .join("")
106 })
107 .unwrap_or_default(),
108 })
109 })
110 .collect::<Vec<_>>()
111 .join("\n"),
112 other => other.as_str().unwrap_or_default().to_owned(),
113 }
114 }
115
116 pub fn ignored_fields(&self) -> Vec<String> {
122 RESERVED
123 .iter()
124 .filter(|f| self.rest.contains_key(**f))
125 .map(|f| (*f).to_owned())
126 .collect()
127 }
128
129 pub fn harness_id(&self) -> Option<&str> {
131 self.metadata
132 .as_ref()?
133 .get("harness_id")
134 .and_then(Value::as_str)
135 }
136}
137
138#[derive(Debug, Clone, Serialize)]
140pub struct OutputItem {
141 #[serde(rename = "type")]
143 pub kind: &'static str,
144 pub id: String,
146 pub status: &'static str,
148 pub role: &'static str,
150 pub content: Vec<ContentPart>,
152}
153
154#[derive(Debug, Clone, Serialize)]
156pub struct ContentPart {
157 #[serde(rename = "type")]
159 pub kind: &'static str,
160 pub text: String,
162}
163
164#[derive(Debug, Clone, Copy, Serialize)]
166pub struct Usage {
167 pub input_tokens: u32,
169 pub output_tokens: u32,
171 pub total_tokens: u32,
173}
174
175#[derive(Debug, Clone, Serialize)]
177pub struct Response {
178 pub id: String,
180 pub object: &'static str,
182 pub created_at: u64,
184 pub status: &'static str,
186 pub error: Option<super::error::UhpError>,
188 pub previous_response_id: Option<String>,
190 pub model: String,
192 pub output: Vec<OutputItem>,
194 pub store: bool,
196 pub usage: Option<Usage>,
198 pub metadata: Map<String, Value>,
200}
201
202#[derive(Debug, Default)]
207pub struct ResponseStore {
208 responses: RwLock<HashMap<String, Response>>,
209 sessions: RwLock<HashMap<String, Vec<String>>>,
210 cancelled: RwLock<HashSet<String>>,
211}
212
213impl ResponseStore {
214 pub fn put(&self, session_id: &str, response: Response) {
216 if let Ok(mut sessions) = self.sessions.write() {
217 sessions
218 .entry(session_id.to_owned())
219 .or_default()
220 .push(response.id.clone());
221 }
222 if let Ok(mut responses) = self.responses.write() {
223 responses.insert(response.id.clone(), response);
224 }
225 }
226
227 pub fn get(&self, id: &str) -> Option<Response> {
229 self.responses.read().ok()?.get(id).cloned()
230 }
231
232 pub fn session_ids(&self) -> Vec<String> {
234 self.sessions
235 .read()
236 .map(|s| s.keys().cloned().collect())
237 .unwrap_or_default()
238 }
239
240 pub fn session_turns(&self, session_id: &str) -> Option<Vec<Response>> {
242 let ids = self.sessions.read().ok()?.get(session_id)?.clone();
243 let responses = self.responses.read().ok()?;
244 Some(
245 ids.iter()
246 .filter_map(|i| responses.get(i).cloned())
247 .collect(),
248 )
249 }
250
251 pub fn cancel(&self, id: &str) -> bool {
258 let running = self.get(id).is_some_and(|r| r.status == "in_progress");
259 if running {
260 if let Ok(mut cancelled) = self.cancelled.write() {
261 cancelled.insert(id.to_owned());
262 }
263 if let Ok(mut responses) = self.responses.write() {
264 if let Some(r) = responses.get_mut(id) {
265 r.status = "cancelled";
266 }
267 }
268 }
269 running
270 }
271
272 pub fn is_cancelled(&self, id: &str) -> bool {
274 self.cancelled.read().is_ok_and(|c| c.contains(id))
275 }
276
277 pub fn settle(&self, response: Response) {
282 if self.is_cancelled(&response.id) {
283 return;
284 }
285 if let Ok(mut responses) = self.responses.write() {
286 responses.insert(response.id.clone(), response);
287 }
288 }
289
290 pub fn running_in(&self, session_id: &str) -> Vec<String> {
292 self.session_turns(session_id)
293 .unwrap_or_default()
294 .into_iter()
295 .filter(|r| r.status == "in_progress")
296 .map(|r| r.id)
297 .collect()
298 }
299
300 pub fn delete_session(&self, session_id: &str) -> bool {
305 let ids = self
306 .sessions
307 .write()
308 .map_or(None, |mut sessions| sessions.remove(session_id));
309 let Some(ids) = ids else { return false };
310 if let Ok(mut responses) = self.responses.write() {
311 for id in &ids {
312 responses.remove(id);
313 }
314 }
315 true
316 }
317
318 pub fn session_of(&self, response_id: &str) -> Option<String> {
320 self.get(response_id)?
321 .metadata
322 .get("session_id")
323 .and_then(Value::as_str)
324 .map(str::to_owned)
325 }
326}
327
328fn now_secs() -> u64 {
330 SystemTime::now()
331 .duration_since(UNIX_EPOCH)
332 .map(|d| d.as_secs())
333 .unwrap_or_default()
334}
335
336const SESSION_PREFIX: &str = "sess_";
338
339fn id_with(prefix: &str) -> String {
341 format!("{prefix}{}", Uuid::new_v4().simple())
342}
343
344#[must_use]
351pub fn is_session_id(id: &str) -> bool {
352 id.strip_prefix(SESSION_PREFIX).is_some_and(|hex| {
353 hex.len() == Simple::LENGTH
354 && hex
355 .bytes()
356 .all(|b| b.is_ascii_digit() || (b'a'..=b'f').contains(&b))
357 })
358}
359
360pub async fn create(
367 State(state): State<UhpState>,
368 Json(request): Json<CreateResponse>,
369) -> Result<AxumResponse, UhpFailure> {
370 if request.background {
371 return Ok(Json(start_background(state, request)).into_response());
372 }
373 if request.stream {
374 let prepared = prepare(&state, &request).await?;
375 let previous = request.previous_response_id.clone();
376 return Ok(super::streaming::sse(state, prepared, previous).into_response());
377 }
378 let outcome = run(&state, &request).await?;
379 Ok(Json(outcome.response).into_response())
380}
381
382fn start_background(state: UhpState, request: CreateResponse) -> Response {
393 let id = id_with("resp_");
394 let session_id = request
395 .previous_response_id
396 .as_ref()
397 .and_then(|prev| state.store.session_of(prev))
398 .unwrap_or_else(|| id_with(SESSION_PREFIX));
399
400 let mut metadata = request.metadata.clone().unwrap_or_default();
401 metadata.insert("session_id".to_owned(), Value::String(session_id.clone()));
402
403 let pending = Response {
404 id: id.clone(),
405 object: "response",
406 created_at: now_secs(),
407 status: "in_progress",
408 error: None,
409 previous_response_id: request.previous_response_id.clone(),
410 model: request.model.clone().unwrap_or_default(),
411 output: Vec::new(),
412 store: true,
413 usage: None,
414 metadata,
415 };
416 state.store.put(&session_id, pending.clone());
417
418 let spawned = pending.clone();
419 tokio::spawn(async move {
420 let finished = match run(&state, &request).await {
421 Ok(outcome) => {
422 let mut r = outcome.response;
425 r.id = id;
426 r
427 }
428 Err(failure) => {
429 let mut r = spawned;
430 r.status = "failed";
431 r.error = Some(failure.into_error());
432 r
433 }
434 };
435 state.store.settle(finished);
436 });
437
438 pending
439}
440
441pub async fn cancel(
447 State(state): State<UhpState>,
448 Path(id): Path<String>,
449) -> Result<Json<Response>, UhpFailure> {
450 if state.store.get(&id).is_none() {
451 return Err(UhpFailure::not_found(
452 "response_not_found",
453 "no response with that id exists",
454 ));
455 }
456 state.store.cancel(&id);
457 state.store.get(&id).map(Json).ok_or_else(|| {
458 UhpFailure::not_found("response_not_found", "no response with that id exists")
459 })
460}
461
462pub struct Outcome {
464 pub response: Response,
466}
467
468pub struct Prepared {
470 pub runner: Arc<dyn LlmProvider>,
472 pub chat: ChatRequest,
474 pub model: String,
476 pub session_id: String,
478 pub id: String,
480 pub metadata: Map<String, Value>,
482 pub requested_model: Option<String>,
484 pub skills: Option<StagedSkills>,
487}
488
489async fn runner_in(
496 provider: CliRunnerType,
497 workdir: &StdPath,
498 harness_args: Vec<String>,
499) -> Result<Arc<dyn LlmProvider>, RunnerError> {
500 fs::create_dir_all(workdir)
501 .map_err(|e| RunnerError::internal(format!("session working folder: {e}")))?;
502
503 let env_override = env::var(provider.env_override_key()).ok();
504 let binary = resolve_binary(provider.binary_name(), env_override.as_deref())?;
505 let mut config = RunnerConfig::new(binary).with_working_directory(workdir.to_path_buf());
506 config.extra_args = write_permission_args(provider);
507 config.extra_args.extend(harness_args);
508
509 Ok(Arc::from(
510 create_runner_with_config(provider, config).await?,
511 ))
512}
513
514fn write_permission_args(provider: CliRunnerType) -> Vec<String> {
535 match provider {
536 CliRunnerType::ClaudeCode => vec![
537 "--permission-mode".to_owned(),
538 "acceptEdits".to_owned(),
539 "--allowed-tools".to_owned(),
540 "Write Edit".to_owned(),
541 ],
542 _ => Vec::new(),
543 }
544}
545
546pub async fn prepare(state: &UhpState, request: &CreateResponse) -> Result<Prepared, UhpFailure> {
561 let chosen = match request.harness_id() {
562 Some(id) => harnesses::resolve(state, id).await.ok_or_else(|| {
563 UhpFailure::not_found("harness_not_found", "no harness with that id is configured")
564 })?,
565 None => harnesses::discover(&state.shared)
566 .await
567 .into_iter()
568 .next()
569 .ok_or_else(|| {
570 UhpFailure::new(
571 ErrorType::InvalidRequestError,
572 "harness_not_found",
573 "this server has no installed harness to run",
574 )
575 })?,
576 };
577 let id = id_with("resp_");
578 let applied = invocation::apply(chosen.provider, &chosen.harness, &id)?;
579
580 let session_id = request
586 .previous_response_id
587 .as_ref()
588 .and_then(|prev| state.store.session_of(prev))
589 .unwrap_or_else(|| id_with(SESSION_PREFIX));
590 let workdir = files::session_workdir(&session_id);
591 let runner = runner_in(chosen.provider, &workdir, applied.args)
592 .await
593 .map_err(|_| {
594 UhpFailure::new(
595 ErrorType::HarnessError,
596 "harness_unavailable",
597 "the harness could not be started",
598 )
599 })?;
600
601 let requested_model = request.model.clone();
602 let model = requested_model
603 .clone()
604 .or_else(|| chosen.harness.default_model.clone())
605 .unwrap_or_else(|| runner.default_model().to_owned());
606
607 let mut chat = ChatRequest::new(vec![ChatMessage::user(request.prompt())])
608 .with_mcp_servers(applied.mcp_servers);
609 chat.model = Some(model.clone());
610
611 let mut metadata = request.metadata.clone().unwrap_or_default();
612 metadata.insert("session_id".to_owned(), Value::String(session_id.clone()));
613 let ignored = request.ignored_fields();
614 if !ignored.is_empty() {
615 metadata.insert(
616 "ignored_fields".to_owned(),
617 Value::Array(ignored.into_iter().map(Value::String).collect()),
618 );
619 }
620
621 Ok(Prepared {
622 runner,
623 chat,
624 model,
625 session_id,
626 id,
627 metadata,
628 requested_model,
629 skills: applied.skills,
630 })
631}
632
633#[must_use]
635pub fn pending_response(p: &Prepared, previous: Option<String>) -> Response {
636 Response {
637 id: p.id.clone(),
638 object: "response",
639 created_at: now_secs(),
640 status: "in_progress",
641 error: None,
642 previous_response_id: previous,
643 model: p.model.clone(),
644 output: Vec::new(),
645 store: true,
646 usage: None,
647 metadata: p.metadata.clone(),
648 }
649}
650
651#[must_use]
653pub fn completed_response(
654 p: Prepared,
655 previous: Option<String>,
656 text: String,
657 usage: Option<Usage>,
658 ran_model: String,
659) -> Response {
660 let mut metadata = p.metadata;
661 if p.requested_model.as_ref().is_some_and(|m| *m != ran_model) {
662 metadata.insert(
663 "requested_model".to_owned(),
664 Value::String(p.requested_model.unwrap_or_default()),
665 );
666 metadata.insert("model_fallback".to_owned(), Value::Bool(true));
667 }
668 Response {
669 id: p.id,
670 object: "response",
671 created_at: now_secs(),
672 status: "completed",
673 error: None,
674 previous_response_id: previous,
675 model: ran_model,
676 output: vec![OutputItem {
677 kind: "message",
678 id: id_with("msg_"),
679 status: "completed",
680 role: "assistant",
681 content: vec![ContentPart {
682 kind: "output_text",
683 text,
684 }],
685 }],
686 store: true,
687 usage,
688 metadata,
689 }
690}
691
692pub async fn run(state: &UhpState, request: &CreateResponse) -> Result<Outcome, UhpFailure> {
698 let prepared = prepare(state, request).await?;
699 let session_id = prepared.session_id.clone();
700 let previous = request.previous_response_id.clone();
701 let model = prepared.model.clone();
702
703 let response = match prepared.runner.complete(&prepared.chat).await {
704 Ok(reply) => {
705 let ran = if reply.model.is_empty() {
706 model
707 } else {
708 reply.model
709 };
710 let usage = reply.usage.map(|u| Usage {
713 input_tokens: u.prompt_tokens,
714 output_tokens: u.completion_tokens,
715 total_tokens: u.total_tokens,
716 });
717 completed_response(prepared, previous, reply.content, usage, ran)
718 }
719 Err(e) => {
720 let mut failed = pending_response(&prepared, previous);
721 failed.status = "failed";
722 failed.error = Some(super::error::UhpError {
723 kind: ErrorType::HarnessError,
724 code: "harness_error".to_owned(),
725 message: e.message,
726 detail: None,
727 });
728 failed
729 }
730 };
731
732 state.store.put(&session_id, response.clone());
733 Ok(Outcome { response })
734}
735
736pub async fn get(
742 State(state): State<UhpState>,
743 Path(id): Path<String>,
744) -> Result<Json<Response>, UhpFailure> {
745 state.store.get(&id).map(Json).ok_or_else(|| {
746 UhpFailure::not_found("response_not_found", "no response with that id exists")
747 })
748}
749
750#[cfg(test)]
751mod tests {
752 use super::*;
753 use serde_json::json;
754
755 fn parse(body: Value) -> CreateResponse {
756 serde_json::from_value(body).expect("request parses") }
758
759 #[test]
760 fn a_bare_string_input_is_one_user_message() {
761 assert_eq!(parse(json!({"input": "hello"})).prompt(), "hello");
762 }
763
764 #[test]
765 fn an_item_array_input_is_flattened_to_its_text() {
766 let r = parse(json!({"input": [
767 {"role": "user", "content": [{"type": "input_text", "text": "first"}]},
768 {"role": "user", "content": "second"}
769 ]}));
770 assert_eq!(r.prompt(), "first\nsecond");
771 }
772
773 #[test]
774 fn ignored_fields_reports_only_what_the_request_carried() {
775 let both = parse(json!({"input": "x", "tools": [], "include": []}));
778 assert_eq!(both.ignored_fields(), vec!["tools", "include"]);
779
780 let one = parse(json!({"input": "x", "include": []}));
781 assert_eq!(one.ignored_fields(), vec!["include"]);
782
783 let neither = parse(json!({"input": "x"}));
784 assert!(
785 neither.ignored_fields().is_empty(),
786 "a request that sent no reserved field must not be told one was ignored"
787 );
788 }
789
790 #[test]
791 fn a_reserved_field_does_not_make_the_request_unparseable() {
792 let r = parse(json!({
795 "input": "x",
796 "tools": [{"type": "function", "name": "probe"}],
797 "metadata": {"harness_id": "chrn_copilot"}
798 }));
799 assert_eq!(r.harness_id(), Some("chrn_copilot"));
800 assert_eq!(r.prompt(), "x");
801 }
802
803 #[test]
804 fn ids_carry_the_prefix_the_object_model_requires() {
805 assert!(id_with("resp_").starts_with("resp_"));
806 assert!(id_with(SESSION_PREFIX).starts_with("sess_"));
807 }
808
809 #[test]
810 fn only_the_session_id_shape_this_server_mints_is_accepted() {
811 assert!(is_session_id(&id_with(SESSION_PREFIX)));
812 for refused in [
813 "..",
814 ".",
815 "",
816 "sess_",
817 "/etc",
818 "sess_0123456789abcdef0123456789abcdef/..",
819 "sess_0123456789abcdef0123456789abcde/",
820 "sess_0123456789ABCDEF0123456789ABCDEF",
821 "sess_0123456789abcdef0123456789abcdef0",
822 "resp_0123456789abcdef0123456789abcdef",
823 ] {
824 assert!(!is_session_id(refused), "{refused:?} is not a minted id");
825 }
826 }
827
828 #[test]
829 fn a_terminal_response_reads_back_unchanged() {
830 let store = ResponseStore::default();
831 let resp = Response {
832 id: "resp_abc".to_owned(),
833 object: "response",
834 created_at: 1,
835 status: "completed",
836 error: None,
837 previous_response_id: None,
838 model: "m".to_owned(),
839 output: Vec::new(),
840 store: true,
841 usage: None,
842 metadata: Map::new(),
843 };
844 store.put("sess_1", resp);
845
846 let got = store.get("resp_abc").expect("stored"); assert_eq!(got.status, "completed", "a terminal state must not change");
848 assert_eq!(store.session_turns("sess_1").map(|t| t.len()), Some(1));
849 assert!(store.get("resp_missing").is_none());
850 }
851}