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, TokenUsage};
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
175impl From<&TokenUsage> for Usage {
176 fn from(usage: &TokenUsage) -> Self {
177 let output_tokens = usage
180 .completion_tokens
181 .saturating_add(usage.reasoning_tokens.unwrap_or(0));
182 let total_tokens = usage
185 .total_tokens
186 .max(usage.prompt_tokens.saturating_add(output_tokens));
187 Self {
188 input_tokens: usage.prompt_tokens,
189 output_tokens,
190 total_tokens,
191 }
192 }
193}
194
195#[derive(Debug, Clone, Serialize)]
197pub struct Response {
198 pub id: String,
200 pub object: &'static str,
202 pub created_at: u64,
204 pub status: &'static str,
206 pub error: Option<super::error::UhpError>,
208 pub previous_response_id: Option<String>,
210 pub model: String,
212 pub output: Vec<OutputItem>,
214 pub store: bool,
216 pub usage: Option<Usage>,
218 pub metadata: Map<String, Value>,
220}
221
222#[derive(Debug, Default)]
227pub struct ResponseStore {
228 responses: RwLock<HashMap<String, Response>>,
229 sessions: RwLock<HashMap<String, Vec<String>>>,
230 cancelled: RwLock<HashSet<String>>,
231}
232
233impl ResponseStore {
234 pub fn put(&self, session_id: &str, response: Response) {
236 if let Ok(mut sessions) = self.sessions.write() {
237 sessions
238 .entry(session_id.to_owned())
239 .or_default()
240 .push(response.id.clone());
241 }
242 if let Ok(mut responses) = self.responses.write() {
243 responses.insert(response.id.clone(), response);
244 }
245 }
246
247 pub fn get(&self, id: &str) -> Option<Response> {
249 self.responses.read().ok()?.get(id).cloned()
250 }
251
252 pub fn session_ids(&self) -> Vec<String> {
254 self.sessions
255 .read()
256 .map(|s| s.keys().cloned().collect())
257 .unwrap_or_default()
258 }
259
260 pub fn session_turns(&self, session_id: &str) -> Option<Vec<Response>> {
262 let ids = self.sessions.read().ok()?.get(session_id)?.clone();
263 let responses = self.responses.read().ok()?;
264 Some(
265 ids.iter()
266 .filter_map(|i| responses.get(i).cloned())
267 .collect(),
268 )
269 }
270
271 pub fn cancel(&self, id: &str) -> bool {
278 let running = self.get(id).is_some_and(|r| r.status == "in_progress");
279 if running {
280 if let Ok(mut cancelled) = self.cancelled.write() {
281 cancelled.insert(id.to_owned());
282 }
283 if let Ok(mut responses) = self.responses.write() {
284 if let Some(r) = responses.get_mut(id) {
285 r.status = "cancelled";
286 }
287 }
288 }
289 running
290 }
291
292 pub fn is_cancelled(&self, id: &str) -> bool {
294 self.cancelled.read().is_ok_and(|c| c.contains(id))
295 }
296
297 pub fn settle(&self, response: Response) {
302 if self.is_cancelled(&response.id) {
303 return;
304 }
305 if let Ok(mut responses) = self.responses.write() {
306 responses.insert(response.id.clone(), response);
307 }
308 }
309
310 pub fn running_in(&self, session_id: &str) -> Vec<String> {
312 self.session_turns(session_id)
313 .unwrap_or_default()
314 .into_iter()
315 .filter(|r| r.status == "in_progress")
316 .map(|r| r.id)
317 .collect()
318 }
319
320 pub fn delete_session(&self, session_id: &str) -> bool {
325 let ids = self
326 .sessions
327 .write()
328 .map_or(None, |mut sessions| sessions.remove(session_id));
329 let Some(ids) = ids else { return false };
330 if let Ok(mut responses) = self.responses.write() {
331 for id in &ids {
332 responses.remove(id);
333 }
334 }
335 true
336 }
337
338 pub fn session_of(&self, response_id: &str) -> Option<String> {
340 self.get(response_id)?
341 .metadata
342 .get("session_id")
343 .and_then(Value::as_str)
344 .map(str::to_owned)
345 }
346}
347
348fn now_secs() -> u64 {
350 SystemTime::now()
351 .duration_since(UNIX_EPOCH)
352 .map(|d| d.as_secs())
353 .unwrap_or_default()
354}
355
356const SESSION_PREFIX: &str = "sess_";
358
359fn id_with(prefix: &str) -> String {
361 format!("{prefix}{}", Uuid::new_v4().simple())
362}
363
364#[must_use]
371pub fn is_session_id(id: &str) -> bool {
372 id.strip_prefix(SESSION_PREFIX).is_some_and(|hex| {
373 hex.len() == Simple::LENGTH
374 && hex
375 .bytes()
376 .all(|b| b.is_ascii_digit() || (b'a'..=b'f').contains(&b))
377 })
378}
379
380pub async fn create(
387 State(state): State<UhpState>,
388 Json(request): Json<CreateResponse>,
389) -> Result<AxumResponse, UhpFailure> {
390 if request.background {
391 return Ok(Json(start_background(state, request)).into_response());
392 }
393 if request.stream {
394 let prepared = prepare(&state, &request).await?;
395 let previous = request.previous_response_id.clone();
396 return Ok(super::streaming::sse(state, prepared, previous).into_response());
397 }
398 let outcome = run(&state, &request).await?;
399 Ok(Json(outcome.response).into_response())
400}
401
402fn start_background(state: UhpState, request: CreateResponse) -> Response {
413 let id = id_with("resp_");
414 let session_id = request
415 .previous_response_id
416 .as_ref()
417 .and_then(|prev| state.store.session_of(prev))
418 .unwrap_or_else(|| id_with(SESSION_PREFIX));
419
420 let mut metadata = request.metadata.clone().unwrap_or_default();
421 metadata.insert("session_id".to_owned(), Value::String(session_id.clone()));
422
423 let pending = Response {
424 id: id.clone(),
425 object: "response",
426 created_at: now_secs(),
427 status: "in_progress",
428 error: None,
429 previous_response_id: request.previous_response_id.clone(),
430 model: request.model.clone().unwrap_or_default(),
431 output: Vec::new(),
432 store: true,
433 usage: None,
434 metadata,
435 };
436 state.store.put(&session_id, pending.clone());
437
438 let spawned = pending.clone();
439 tokio::spawn(async move {
440 let finished = match run(&state, &request).await {
441 Ok(outcome) => {
442 let mut r = outcome.response;
445 r.id = id;
446 r
447 }
448 Err(failure) => {
449 let mut r = spawned;
450 r.status = "failed";
451 r.error = Some(failure.into_error());
452 r
453 }
454 };
455 state.store.settle(finished);
456 });
457
458 pending
459}
460
461pub async fn cancel(
467 State(state): State<UhpState>,
468 Path(id): Path<String>,
469) -> Result<Json<Response>, UhpFailure> {
470 if state.store.get(&id).is_none() {
471 return Err(UhpFailure::not_found(
472 "response_not_found",
473 "no response with that id exists",
474 ));
475 }
476 state.store.cancel(&id);
477 state.store.get(&id).map(Json).ok_or_else(|| {
478 UhpFailure::not_found("response_not_found", "no response with that id exists")
479 })
480}
481
482pub struct Outcome {
484 pub response: Response,
486}
487
488pub struct Prepared {
490 pub runner: Arc<dyn LlmProvider>,
492 pub chat: ChatRequest,
494 pub model: String,
496 pub session_id: String,
498 pub id: String,
500 pub metadata: Map<String, Value>,
502 pub requested_model: Option<String>,
504 pub skills: Option<StagedSkills>,
507}
508
509async fn runner_in(
516 provider: CliRunnerType,
517 workdir: &StdPath,
518 harness_args: Vec<String>,
519) -> Result<Arc<dyn LlmProvider>, RunnerError> {
520 fs::create_dir_all(workdir)
521 .map_err(|e| RunnerError::internal(format!("session working folder: {e}")))?;
522
523 let env_override = env::var(provider.env_override_key()).ok();
524 let binary = resolve_binary(provider.binary_name(), env_override.as_deref())?;
525 let mut config = RunnerConfig::new(binary).with_working_directory(workdir.to_path_buf());
526 config.extra_args = write_permission_args(provider);
527 config.extra_args.extend(harness_args);
528
529 Ok(Arc::from(
530 create_runner_with_config(provider, config).await?,
531 ))
532}
533
534fn write_permission_args(provider: CliRunnerType) -> Vec<String> {
555 match provider {
556 CliRunnerType::ClaudeCode => vec![
557 "--permission-mode".to_owned(),
558 "acceptEdits".to_owned(),
559 "--allowed-tools".to_owned(),
560 "Write Edit".to_owned(),
561 ],
562 _ => Vec::new(),
563 }
564}
565
566pub async fn prepare(state: &UhpState, request: &CreateResponse) -> Result<Prepared, UhpFailure> {
581 let chosen = match request.harness_id() {
582 Some(id) => harnesses::resolve(state, id).await.ok_or_else(|| {
583 UhpFailure::not_found("harness_not_found", "no harness with that id is configured")
584 })?,
585 None => harnesses::discover(&state.shared)
586 .await
587 .into_iter()
588 .next()
589 .ok_or_else(|| {
590 UhpFailure::new(
591 ErrorType::InvalidRequestError,
592 "harness_not_found",
593 "this server has no installed harness to run",
594 )
595 })?,
596 };
597 let id = id_with("resp_");
598 let applied = invocation::apply(chosen.provider, &chosen.harness, &id)?;
599
600 let session_id = request
606 .previous_response_id
607 .as_ref()
608 .and_then(|prev| state.store.session_of(prev))
609 .unwrap_or_else(|| id_with(SESSION_PREFIX));
610 let workdir = files::session_workdir(&session_id);
611 let runner = runner_in(chosen.provider, &workdir, applied.args)
612 .await
613 .map_err(|_| {
614 UhpFailure::new(
615 ErrorType::HarnessError,
616 "harness_unavailable",
617 "the harness could not be started",
618 )
619 })?;
620
621 let requested_model = request.model.clone();
622 let model = requested_model
623 .clone()
624 .or_else(|| chosen.harness.default_model.clone())
625 .unwrap_or_else(|| runner.default_model().to_owned());
626
627 let mut chat = ChatRequest::new(vec![ChatMessage::user(request.prompt())])
628 .with_mcp_servers(applied.mcp_servers);
629 chat.model = Some(model.clone());
630
631 let mut metadata = request.metadata.clone().unwrap_or_default();
632 metadata.insert("session_id".to_owned(), Value::String(session_id.clone()));
633 let ignored = request.ignored_fields();
634 if !ignored.is_empty() {
635 metadata.insert(
636 "ignored_fields".to_owned(),
637 Value::Array(ignored.into_iter().map(Value::String).collect()),
638 );
639 }
640
641 Ok(Prepared {
642 runner,
643 chat,
644 model,
645 session_id,
646 id,
647 metadata,
648 requested_model,
649 skills: applied.skills,
650 })
651}
652
653#[must_use]
655pub fn pending_response(p: &Prepared, previous: Option<String>) -> Response {
656 Response {
657 id: p.id.clone(),
658 object: "response",
659 created_at: now_secs(),
660 status: "in_progress",
661 error: None,
662 previous_response_id: previous,
663 model: p.model.clone(),
664 output: Vec::new(),
665 store: true,
666 usage: None,
667 metadata: p.metadata.clone(),
668 }
669}
670
671#[must_use]
673pub fn completed_response(
674 p: Prepared,
675 previous: Option<String>,
676 text: String,
677 usage: Option<Usage>,
678 ran_model: String,
679) -> Response {
680 let mut metadata = p.metadata;
681 if p.requested_model.as_ref().is_some_and(|m| *m != ran_model) {
682 metadata.insert(
683 "requested_model".to_owned(),
684 Value::String(p.requested_model.unwrap_or_default()),
685 );
686 metadata.insert("model_fallback".to_owned(), Value::Bool(true));
687 }
688 Response {
689 id: p.id,
690 object: "response",
691 created_at: now_secs(),
692 status: "completed",
693 error: None,
694 previous_response_id: previous,
695 model: ran_model,
696 output: vec![OutputItem {
697 kind: "message",
698 id: id_with("msg_"),
699 status: "completed",
700 role: "assistant",
701 content: vec![ContentPart {
702 kind: "output_text",
703 text,
704 }],
705 }],
706 store: true,
707 usage,
708 metadata,
709 }
710}
711
712pub async fn run(state: &UhpState, request: &CreateResponse) -> Result<Outcome, UhpFailure> {
718 let prepared = prepare(state, request).await?;
719 let session_id = prepared.session_id.clone();
720 let previous = request.previous_response_id.clone();
721 let model = prepared.model.clone();
722
723 let response = match prepared.runner.complete(&prepared.chat).await {
724 Ok(reply) => {
725 let ran = if reply.model.is_empty() {
726 model
727 } else {
728 reply.model
729 };
730 let usage = reply.usage.as_ref().map(Usage::from);
733 completed_response(prepared, previous, reply.content, usage, ran)
734 }
735 Err(e) => {
736 let mut failed = pending_response(&prepared, previous);
737 failed.status = "failed";
738 failed.error = Some(super::error::UhpError {
739 kind: ErrorType::HarnessError,
740 code: "harness_error".to_owned(),
741 message: e.message,
742 detail: None,
743 });
744 failed
745 }
746 };
747
748 state.store.put(&session_id, response.clone());
749 Ok(Outcome { response })
750}
751
752pub async fn get(
758 State(state): State<UhpState>,
759 Path(id): Path<String>,
760) -> Result<Json<Response>, UhpFailure> {
761 state.store.get(&id).map(Json).ok_or_else(|| {
762 UhpFailure::not_found("response_not_found", "no response with that id exists")
763 })
764}
765
766#[cfg(test)]
767mod tests {
768 use super::*;
769 use serde_json::json;
770
771 fn parse(body: Value) -> CreateResponse {
772 serde_json::from_value(body).expect("request parses") }
774
775 #[test]
776 fn a_bare_string_input_is_one_user_message() {
777 assert_eq!(parse(json!({"input": "hello"})).prompt(), "hello");
778 }
779
780 #[test]
781 fn an_item_array_input_is_flattened_to_its_text() {
782 let r = parse(json!({"input": [
783 {"role": "user", "content": [{"type": "input_text", "text": "first"}]},
784 {"role": "user", "content": "second"}
785 ]}));
786 assert_eq!(r.prompt(), "first\nsecond");
787 }
788
789 #[test]
790 fn ignored_fields_reports_only_what_the_request_carried() {
791 let both = parse(json!({"input": "x", "tools": [], "include": []}));
794 assert_eq!(both.ignored_fields(), vec!["tools", "include"]);
795
796 let one = parse(json!({"input": "x", "include": []}));
797 assert_eq!(one.ignored_fields(), vec!["include"]);
798
799 let neither = parse(json!({"input": "x"}));
800 assert!(
801 neither.ignored_fields().is_empty(),
802 "a request that sent no reserved field must not be told one was ignored"
803 );
804 }
805
806 #[test]
807 fn a_reserved_field_does_not_make_the_request_unparseable() {
808 let r = parse(json!({
811 "input": "x",
812 "tools": [{"type": "function", "name": "probe"}],
813 "metadata": {"harness_id": "chrn_copilot"}
814 }));
815 assert_eq!(r.harness_id(), Some("chrn_copilot"));
816 assert_eq!(r.prompt(), "x");
817 }
818
819 #[test]
820 fn ids_carry_the_prefix_the_object_model_requires() {
821 assert!(id_with("resp_").starts_with("resp_"));
822 assert!(id_with(SESSION_PREFIX).starts_with("sess_"));
823 }
824
825 #[test]
826 fn only_the_session_id_shape_this_server_mints_is_accepted() {
827 assert!(is_session_id(&id_with(SESSION_PREFIX)));
828 for refused in [
829 "..",
830 ".",
831 "",
832 "sess_",
833 "/etc",
834 "sess_0123456789abcdef0123456789abcdef/..",
835 "sess_0123456789abcdef0123456789abcde/",
836 "sess_0123456789ABCDEF0123456789ABCDEF",
837 "sess_0123456789abcdef0123456789abcdef0",
838 "resp_0123456789abcdef0123456789abcdef",
839 ] {
840 assert!(!is_session_id(refused), "{refused:?} is not a minted id");
841 }
842 }
843
844 #[test]
845 fn a_terminal_response_reads_back_unchanged() {
846 let store = ResponseStore::default();
847 let resp = Response {
848 id: "resp_abc".to_owned(),
849 object: "response",
850 created_at: 1,
851 status: "completed",
852 error: None,
853 previous_response_id: None,
854 model: "m".to_owned(),
855 output: Vec::new(),
856 store: true,
857 usage: None,
858 metadata: Map::new(),
859 };
860 store.put("sess_1", resp);
861
862 let got = store.get("resp_abc").expect("stored"); assert_eq!(got.status, "completed", "a terminal state must not change");
864 assert_eq!(store.session_turns("sess_1").map(|t| t.len()), Some(1));
865 assert!(store.get("resp_missing").is_none());
866 }
867
868 #[test]
869 fn output_counts_reasoning_and_the_total_covers_both_sides() {
870 let usage = Usage::from(&TokenUsage::new(100, 7, 107).with_reasoning(Some(20)));
872 assert_eq!(usage.input_tokens, 100);
873 assert_eq!(usage.output_tokens, 27);
874 assert_eq!(usage.total_tokens, 127);
875
876 let usage = Usage::from(&TokenUsage::new(100, 7, 127).with_reasoning(Some(20)));
878 assert_eq!(usage.output_tokens, 27);
879 assert_eq!(usage.total_tokens, 127);
880
881 let usage = Usage::from(&TokenUsage::new(10, 5, 15));
883 assert_eq!((usage.output_tokens, usage.total_tokens), (5, 15));
884 }
885}