1use std::cmp::Reverse;
19use std::collections::HashSet;
20use std::io::{BufRead, IsTerminal, Write};
21use std::sync::{Arc, Mutex};
22
23use rpi_agent::events::AgentEvent;
24use rpi_ai::types::{AssistantMessage, Content, ImageContent, StopReason, UserMessage};
25use rpi_harness::agent_harness::{
26 AgentHarness, AgentLane, CompactionOutcome, CompactionResult, HarnessRunOutcome,
27};
28use rpi_harness::events::{HarnessEvent, RunEndOutcome};
29use rpi_harness::session::session::SessionEntryCommit;
30use rpi_harness::session::types::{
31 BranchBounds, Entry, EntryCursor, EntryOrder, EntryQuery, LaneRecord, RecordQuery,
32};
33
34use crate::args::Args;
35use crate::rpc::protocol as rpc_protocol;
36use crate::rpc::writer::{start_rpc_writer, write_agent_event, write_line, RpcWriter};
37
38pub fn assistant_text(msg: &AssistantMessage) -> String {
41 msg.content
42 .iter()
43 .filter_map(|c| match c {
44 Content::Text(t) => Some(t.text.clone()),
45 _ => None,
46 })
47 .collect()
48}
49
50pub fn outcome_exit_code(outcome: &HarnessRunOutcome) -> i32 {
53 match outcome {
54 HarnessRunOutcome::Failed { .. } | HarnessRunOutcome::Aborted { .. } => 1,
55 _ => 0,
56 }
57}
58
59pub async fn print(
63 harness: &AgentHarness,
64 args: &Args,
65 initial: Option<String>,
66 extra_messages: &[String],
67 initial_images: Vec<ImageContent>,
68) -> i32 {
69 let lane: Arc<dyn AgentLane> = harness.lane("main");
70
71 let mut last_exit = 0;
72 let mut persistence_failed = false;
73 let mut last_msg: Option<AssistantMessage> = None;
74
75 let mut prompts: Vec<String> = Vec::new();
78 if let Some(init) = initial {
79 prompts.push(init);
80 }
81 for m in extra_messages {
82 prompts.push(m.clone());
83 }
84
85 if prompts.is_empty() {
86 return 0;
88 }
89
90 let mut images = initial_images;
91 for prompt in prompts {
92 let outcome = lane.prompt_text(&prompt, std::mem::take(&mut images)).await;
93 for failure in lane.take_write_failures() {
97 persistence_failed |= args.experimental_native_store.is_some();
98 eprintln!(
99 "session not saved: could not persist {} ({})",
100 failure.what, failure.message
101 );
102 }
103 match outcome {
104 Ok(result) => {
105 last_exit = outcome_exit_code(&result.outcome);
106 match &result.outcome {
107 HarnessRunOutcome::Completed { final_message, .. }
108 | HarnessRunOutcome::Aborted { final_message, .. } => {
109 last_msg = Some(final_message.clone());
110 }
111 HarnessRunOutcome::Failed {
112 error,
113 final_message,
114 ..
115 } => {
116 if let Some(m) = final_message {
117 if m.stop_reason == StopReason::Error {
118 if let Some(em) = &m.error_message {
119 eprintln!("{em}");
120 }
121 }
122 }
123 if error.code == "SessionStorage" {
124 eprintln!("{}", session_storage_notice(&error.message));
125 } else {
126 eprintln!("run failed: {error:?}");
127 }
128 }
129 HarnessRunOutcome::Suspended { .. } => {
130 eprintln!("run suspended (deferred) — resume is not supported in v1");
131 last_exit = 1;
132 }
133 }
134 }
135 Err(e) => {
136 eprintln!("prompt rejected: {e}");
137 return 1;
138 }
139 }
140 }
141
142 if let Some(m) = &last_msg {
144 match m.stop_reason {
145 StopReason::Error => {
146 if let Some(em) = &m.error_message {
147 eprintln!("{em}");
148 }
149 last_exit = 1;
150 }
151 StopReason::Aborted => {
152 eprintln!("request aborted");
153 last_exit = 1;
154 }
155 _ => {
156 let text = assistant_text(m);
157 let mut out = std::io::stdout();
158 let _ = out.write_all(text.as_bytes());
159 if !text.ends_with('\n') {
160 let _ = out.write_all(b"\n");
161 }
162 let _ = out.flush();
163 }
164 }
165 }
166
167 if persistence_failed {
168 1
169 } else {
170 last_exit
171 }
172}
173
174pub async fn json(
178 harness: &AgentHarness,
179 args: &Args,
180 initial: Option<String>,
181 extra_messages: &[String],
182 initial_images: Vec<ImageContent>,
183 mut agent_events: Option<tokio::sync::broadcast::Receiver<AgentEvent>>,
184) -> i32 {
185 let lane: Arc<dyn AgentLane> = harness.lane("main");
186 let collected: Arc<Mutex<Vec<HarnessEvent>>> = Arc::new(Mutex::new(Vec::new()));
187 let collected_for_watch = collected.clone();
188
189 let mut watch = harness.events().watch(|| ());
192 watch.start(Arc::new(move |event: &HarnessEvent| {
193 emit_json_event(event);
195 collected_for_watch.lock().unwrap().push(event.clone());
196 }));
197 std::mem::forget(watch);
200
201 let (agent_done_tx, mut agent_done_rx) = tokio::sync::mpsc::unbounded_channel::<()>();
204 let agent_event_task = agent_events.take().map(|mut rx| {
205 tokio::spawn(async move {
206 let done_tx = agent_done_tx;
207 loop {
208 match rx.recv().await {
209 Ok(event) => {
210 let terminal = event.is_terminal();
211 emit_agent_event(&event);
212 if terminal {
213 let _ = done_tx.send(());
214 }
215 }
216 Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => continue,
217 Err(tokio::sync::broadcast::error::RecvError::Closed) => break,
218 }
219 }
220 })
221 });
222
223 let mut prompts: Vec<String> = Vec::new();
224 if let Some(init) = initial {
225 prompts.push(init);
226 }
227 for m in extra_messages {
228 prompts.push(m.clone());
229 }
230
231 let mut last_exit = 0;
232 let mut persistence_failed = false;
233 let mut final_outcome: Option<HarnessRunOutcome> = None;
234
235 let mut images = initial_images;
236 for prompt in prompts {
237 match lane.prompt_text(&prompt, std::mem::take(&mut images)).await {
238 Ok(result) => {
239 let _ = tokio::time::timeout(
245 std::time::Duration::from_millis(250),
246 agent_done_rx.recv(),
247 )
248 .await;
249 last_exit = outcome_exit_code(&result.outcome);
250 final_outcome = Some(result.outcome);
251 }
252 Err(e) => {
253 let line = serde_json::json!({
255 "type": "error",
256 "error": e.to_string(),
257 });
258 println!("{line}");
259 if args.experimental_native_store.is_some() {
260 for failure in lane.take_write_failures() {
261 persistence_failed = true;
262 println!(
263 "{}",
264 serde_json::json!({"type":"error","kind":"session_write_failed","error":failure.message,"what":failure.what})
265 );
266 }
267 last_exit = 1;
268 break;
269 }
270 if let Some(task) = agent_event_task {
271 task.abort();
272 }
273 return 1;
274 }
275 }
276 if args.experimental_native_store.is_some() {
277 for failure in lane.take_write_failures() {
278 persistence_failed = true;
279 println!(
280 "{}",
281 serde_json::json!({"type":"error","kind":"session_write_failed","error":failure.message,"what":failure.what})
282 );
283 }
284 }
285 }
286
287 let (outcome_str, final_text) = match final_outcome {
289 Some(HarnessRunOutcome::Completed { final_message, .. }) => {
290 ("completed", Some(assistant_text(&final_message)))
291 }
292 Some(HarnessRunOutcome::Aborted { final_message, .. }) => {
293 ("aborted", Some(assistant_text(&final_message)))
294 }
295 Some(HarnessRunOutcome::Failed { final_message, .. }) => {
296 let t = final_message.as_ref().map(assistant_text);
297 ("failed", t)
298 }
299 Some(HarnessRunOutcome::Suspended { .. }) => ("suspended", None),
300 None if last_exit != 0 => ("failed", None),
301 None => ("idle", None),
302 };
303 let mut result_line = serde_json::json!({
304 "type": "result",
305 "outcome": if persistence_failed{"failed"}else{outcome_str},
306 "finalText": final_text,
307 });
308 if persistence_failed {
309 result_line["persistenceFailed"] = true.into();
310 last_exit = 1;
311 }
312 println!("{result_line}");
313 if let Some(task) = agent_event_task {
314 task.abort();
315 }
316 last_exit
317}
318
319fn emit_agent_event(event: &AgentEvent) {
320 println!("{}", agent_event_json(event));
321}
322
323pub async fn rpc(
342 harness: &AgentHarness,
343 mut agent_events: Option<tokio::sync::broadcast::Receiver<AgentEvent>>,
344 model_catalog: Vec<rpi_ai::Model>,
345 ui_dialog: rpi_extensions::UiDialogMailbox,
346 oauth_snapshot: Option<Arc<rpi_extensions::RegistrySnapshot>>,
347 message_updates: crate::args::RpcMessageUpdates,
348) -> i32 {
349 rpc_with_io(
350 harness,
351 agent_events.take(),
352 model_catalog,
353 ui_dialog,
354 oauth_snapshot,
355 message_updates,
356 tokio::io::BufReader::new(tokio::io::stdin()),
357 std::io::stdout(),
358 )
359 .await
360}
361
362struct RpcPhaseTiming {
365 phase: &'static str,
366 started: Option<std::time::Instant>,
367}
368
369impl RpcPhaseTiming {
370 fn start(enabled: bool, phase: &'static str) -> Self {
371 let started = enabled.then(std::time::Instant::now);
372 if enabled {
373 eprintln!("RPI_RPC_DIAGNOSTICS phase={phase} event=begin");
374 }
375 Self { phase, started }
376 }
377
378 fn finish(self) {
379 if let Some(started) = self.started {
380 eprintln!(
381 "RPI_RPC_DIAGNOSTICS phase={} event=end elapsed_ns={}",
382 self.phase,
383 started.elapsed().as_nanos()
384 );
385 }
386 }
387}
388
389async fn rpc_with_io<R, W>(
390 harness: &AgentHarness,
391 mut agent_events: Option<tokio::sync::broadcast::Receiver<AgentEvent>>,
392 model_catalog: Vec<rpi_ai::Model>,
393 ui_dialog: rpi_extensions::UiDialogMailbox,
394 oauth_snapshot: Option<Arc<rpi_extensions::RegistrySnapshot>>,
395 message_updates: crate::args::RpcMessageUpdates,
396 mut input: R,
397 output: W,
398) -> i32
399where
400 R: tokio::io::AsyncBufRead + Unpin,
401 W: Write + Send + 'static,
402{
403 let phase_diagnostics =
404 std::env::var_os("RPI_RPC_DIAGNOSTICS").is_some_and(|value| value != "0");
405 let setup_timing = RpcPhaseTiming::start(phase_diagnostics, "setup_to_ready_enqueue");
406 let mut command_tasks: Vec<RetainedRpcTask<()>> = Vec::new();
407 let lane: Arc<dyn AgentLane> = harness.lane("main");
408 let (stdout, writer_task, writer_failed) = start_rpc_writer(output, message_updates);
411 let writer_task = RetainedRpcTask::new(writer_task, writer_failed.clone());
412 let pump_stop = tokio_util::sync::CancellationToken::new();
413 let external_tools =
414 crate::rpc::external_tools::ExternalTools::new(stdout.clone(), Arc::downgrade(&lane));
415 let mut external_names = Vec::new();
416 let file_io_router = crate::session::file_io_router(harness);
417 let mut owner = RpcOwner {
418 harness,
419 file_io: None,
420 backend: external_tools.clone(),
421 run_cancel: tokio_util::sync::CancellationToken::new(),
422 pump_stop: pump_stop.clone(),
423 failed: writer_failed.clone(),
424 mailbox: ui_dialog.clone(),
425 armed: true,
426 };
427
428 let (commit_tx, mut commit_rx) = tokio::sync::broadcast::channel::<SessionEntryCommit>(256);
432 harness
433 .session()
434 .set_entry_commit_publisher(Some(commit_tx.clone()));
435 let commit_task = {
436 let out = stdout.clone();
437 tokio::spawn(async move {
438 loop {
439 match commit_rx.recv().await {
440 Ok(commit) => {
441 write_line(
442 &out,
443 &serde_json::json!({
444 "type": "session_entry_committed",
445 "entryId": commit.entry_id,
446 "entryType": commit.entry_type,
447 "customType": commit.custom_type,
448 "lane": commit.lane,
449 }),
450 )
451 .await;
452 }
453 Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => continue,
454 Err(tokio::sync::broadcast::error::RecvError::Closed) => break,
455 }
456 }
457 })
458 };
459
460 let commit_task = RetainedRpcTask::new(commit_task, pump_stop.clone());
461
462 let (event_flush_tx, mut event_flush_rx) =
465 tokio::sync::mpsc::unbounded_channel::<tokio::sync::oneshot::Sender<()>>();
466 let event_task = agent_events.take().map(|mut rx| {
468 let out = stdout.clone();
469 let stop = pump_stop.clone();
470 tokio::spawn(async move {
471 loop {
472 let received = tokio::select! {
473 _ = stop.cancelled() => {
474 drain_rpc_events(&out, &mut rx).await;
475 break;
476 }
477 Some(done) = event_flush_rx.recv() => {
478 drain_rpc_events(&out, &mut rx).await;
479 let _ = done.send(());
480 continue;
481 }
482 event = rx.recv() => event,
483 };
484 match received {
485 Ok(event) => write_agent_event(&out, &event).await,
486 Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => continue,
487 Err(tokio::sync::broadcast::error::RecvError::Closed) => break,
488 }
489 }
490 })
491 });
492
493 let event_task = event_task.map(|task| RetainedRpcTask::new(task, pump_stop.clone()));
494 let event_flush = event_task.as_ref().map(|_| event_flush_tx);
495
496 ui_dialog.attach();
502 let dialog_task = {
503 let mailbox = ui_dialog.clone();
504 let out = stdout.clone();
505 let stop = pump_stop.clone();
506 tokio::spawn(async move {
507 let mut shown = std::collections::HashSet::new();
508 while !stop.is_cancelled() {
509 mailbox.set_external_takeover(true);
512 while let Some(request) =
516 mailbox.take_pending_filtered(|r| !mailbox.is_reserved(&r.request_id))
517 {
518 let event = crate::remote::protocol::RemoteEvent::UiDialogRequest {
519 request_id: request.request_id.clone(),
520 tool_call_id: request.tool_call_id,
521 ui: request.ui,
522 };
523 write_line(&out, &serde_json::to_value(event).unwrap()).await;
524 shown.insert(request.request_id);
525 }
526 let closed: Vec<String> = shown
527 .iter()
528 .filter(|id| !mailbox.is_pending(id))
529 .cloned()
530 .collect();
531 for request_id in closed {
532 shown.remove(&request_id);
533 let event = crate::remote::protocol::RemoteEvent::UiDialogClosed { request_id };
534 write_line(&out, &serde_json::to_value(event).unwrap()).await;
535 }
536 tokio::select! {
537 _ = stop.cancelled() => break,
538 _ = tokio::time::sleep(std::time::Duration::from_millis(30)) => {}
539 }
540 }
541 })
542 };
543
544 let dialog_task = RetainedRpcTask::new(dialog_task, pump_stop.clone());
545
546 let run_active = Arc::new(std::sync::atomic::AtomicBool::new(false));
550 let mut run_task: Option<OwnedRpcRun> = None;
551 let mut run_cancel = tokio_util::sync::CancellationToken::new();
552 let mut cancel_on_exit = false;
553
554 let binding = match rpi_harness::session::SessionValues::load(harness.session()).await {
555 Ok(values) => crate::resource_selection::binding_state(&values),
556 Err(error) => {
557 eprintln!("resource binding read failed: {error}");
558 return 1;
559 }
560 };
561 let mut ready = rpc_protocol::ready(message_updates);
562 ready["resourceState"] = serde_json::json!({
563 "version": 1, "enabled": harness.has_exact_resources(), "command": "get_resource_state",
564 "maxSnapshotBytes": crate::resource_selection::MAX_STATE_BYTES,
565 "changes": "restart-required", "setActiveTools": if harness.has_exact_resources() { "unsupported" } else { "legacy" },
566 "configureExternalTools": if harness.has_exact_resources() { "unsupported" } else { "supported" },
567 "automaticPersistenceRecovery": false, "hotApply": false,
568 "packageScopes": ["project", "global"], "defaultPackageScope": "project",
569 "transition": crate::resource_selection::transition_capability(),
570 "binding": binding
571 });
572 if harness.has_exact_resources() {
573 if let Ok(state) = crate::resource_selection::get_resource_state(harness).await {
574 for key in [
575 "hostId",
576 "host",
577 "sessionId",
578 "selection",
579 "binding",
580 "transition",
581 ] {
582 ready["resourceState"][key] = state[key].clone();
583 }
584 } else {
585 ready["resourceState"]["enabled"] = serde_json::json!(false);
586 }
587 }
588 write_line(&stdout, &ready).await;
589 setup_timing.finish();
590
591 loop {
592 let incoming = tokio::select! {
593 _ = writer_failed.cancelled() => {
594 cancel_on_exit = true;
595 break;
596 }
597 line = crate::rpc::framing::read_frame(&mut input) => line,
598 };
599 let line = match incoming {
600 Ok(Some(line)) => line,
601 Ok(None) => break, Err(error) => {
603 eprintln!("[rpi] rpc stdin error: {error}");
604 run_cancel.cancel();
605 let _ = lane.abort().await;
606 external_tools.shutdown();
607 writer_failed.cancel();
608 cancel_on_exit = true;
609 break;
610 }
611 };
612 let line = match std::str::from_utf8(&line) {
613 Ok(line) => line.trim(),
614 Err(_) => {
615 run_cancel.cancel();
616 let _ = lane.abort().await;
617 external_tools.shutdown();
618 writer_failed.cancel();
619 cancel_on_exit = true;
620 break;
621 }
622 };
623 if line.is_empty() {
624 continue;
625 }
626 let msg: serde_json::Value = match serde_json::from_str(line) {
627 Ok(value) => value,
628 Err(error) => {
629 write_line(
630 &stdout,
631 &rpc_protocol::response(
632 None,
633 "invalid_request",
634 Err(format!("invalid json: {error}")),
635 ),
636 )
637 .await;
638 continue;
639 }
640 };
641 let id = msg.get("id").cloned().unwrap_or(serde_json::Value::Null);
642 let kind = msg
643 .get("type")
644 .and_then(|v| v.as_str())
645 .unwrap_or("")
646 .to_string();
647
648 if kind == "get_resource_state" && serde_json::to_vec(&id).map_or(true, |v| v.len() > 4096)
649 {
650 write_line(
651 &stdout,
652 &rpc_protocol::response(
653 None,
654 &kind,
655 Err("resource state request id exceeds 4096 bytes".into()),
656 ),
657 )
658 .await;
659 continue;
660 }
661 if kind != "prompt"
662 && rpc_protocol::SUPPORTED_COMMANDS.contains(&kind.as_str())
663 && serde_json::from_value::<rpc_protocol::Request>(msg.clone()).is_err()
664 {
665 write_line(
666 &stdout,
667 &rpc_protocol::response(
668 Some(id.clone()),
669 &kind,
670 Err("invalid command payload".to_string()),
671 ),
672 )
673 .await;
674 continue;
675 }
676
677 match kind.as_str() {
678 "ui_dialog_response" => {
679 let request_id = msg
680 .get("requestId")
681 .and_then(|v| v.as_str())
682 .unwrap_or("")
683 .to_string();
684 let result = if msg
685 .get("cancelled")
686 .and_then(|v| v.as_bool())
687 .unwrap_or(false)
688 {
689 ui_dialog.cancel(&request_id)
690 } else {
691 let answer = msg
692 .get("answer")
693 .cloned()
694 .unwrap_or(serde_json::Value::Null);
695 ui_dialog.respond(&request_id, answer)
696 };
697 if msg.get("id").is_some() {
698 let response = match result {
699 Ok(()) => rpc_protocol::response(
700 Some(id),
701 "ui_dialog_response",
702 Ok(Some(
703 serde_json::json!({"accepted":true,"requestId":request_id}),
704 )),
705 ),
706 Err(error) => {
707 rpc_protocol::response(Some(id), "ui_dialog_response", Err(error))
708 }
709 };
710 write_line(&stdout, &response).await;
711 } else if let Err(error) = result {
712 eprintln!("[rpi] ui_dialog_response: {error}");
713 }
714 continue;
715 }
716 "execute_command" => {
717 command_tasks.retain(|task| !task.is_finished());
718 let snapshot = oauth_snapshot.clone();
719 let request = msg.clone();
720 let out = stdout.clone();
721 command_tasks.push(RetainedRpcTask::new(
722 tokio::spawn(async move {
723 let result = tokio::task::spawn_blocking(move || {
724 rpc_execute_command(snapshot.as_deref(), &request)
725 })
726 .await
727 .unwrap_or_else(|_| Err("command task failed".into()));
728 write_line(
729 &out,
730 &rpc_protocol::response(Some(id), "execute_command", result.map(Some)),
731 )
732 .await;
733 }),
734 writer_failed.clone(),
735 ));
736 continue;
737 }
738 "stop" => {
739 cancel_on_exit = true;
740 run_cancel.cancel();
741 let _ = lane.abort().await;
742 external_tools.shutdown();
743 write_line(
744 &stdout,
745 &rpc_protocol::response(
746 Some(id.clone()),
747 "stop",
748 Ok(Some(serde_json::json!({"stopped": true}))),
749 ),
750 )
751 .await;
752 break;
753 }
754 "prompt" | "resume_pending" | "compact" => {
755 let resume = kind == "resume_pending";
756 let compact = kind == "compact";
757 if run_active.swap(true, std::sync::atomic::Ordering::AcqRel) {
758 write_line(
759 &stdout,
760 &rpc_protocol::response(
761 Some(id.clone()),
762 &kind,
763 Err("a run is already active".to_string()),
764 ),
765 )
766 .await;
767 continue;
768 }
769 if let Some(task) = run_task.take() {
770 let _ = task.await;
771 }
772 run_cancel = tokio_util::sync::CancellationToken::new();
773 owner.run_cancel = run_cancel.clone();
774 let cancel_task = run_cancel.clone();
775 if resume && !lane.has_pending_resume().await {
778 run_active.store(false, std::sync::atomic::Ordering::Release);
779 write_line(
780 &stdout,
781 &rpc_protocol::response(
782 Some(id),
783 "resume_pending",
784 Ok(Some(serde_json::json!({"resumed": false}))),
785 ),
786 )
787 .await;
788 continue;
789 }
790 let lane_task = Arc::clone(&lane);
791 let out_task = stdout.clone();
792 let active_task = Arc::clone(&run_active);
793 let mailbox_task = ui_dialog.clone();
794 let flush_events = event_flush.clone();
795 let instructions = msg
796 .get("customInstructions")
797 .and_then(|v| v.as_str())
798 .map(str::to_owned);
799 run_task = Some(OwnedRpcRun::new(
804 tokio::spawn(async move {
805 #[cfg(test)]
806 let _finished =
807 crate::rpc::prompt_content::validation_test_hook::for_command(&msg)
808 .map(crate::rpc::prompt_content::validation_test_hook::RunFinished);
809 let validation_timing =
810 RpcPhaseTiming::start(phase_diagnostics, "validation");
811 let content = if resume || compact {
812 Ok(crate::rpc::prompt_content::ParsedPrompt::Legacy(
813 String::new(),
814 ))
815 } else {
816 rpc_validate_prompt(&lane_task, msg, &cancel_task).await
817 };
818 validation_timing.finish();
819 let content = match content {
820 Ok(content) => content,
821 Err(error) => {
822 write_line(
823 &out_task,
824 &rpc_protocol::response(Some(id), &kind, Err(error)),
825 )
826 .await;
827 active_task.store(false, std::sync::atomic::Ordering::Release);
828 return;
829 }
830 };
831 write_line(
832 &out_task,
833 &rpc_protocol::response(
834 Some(id),
835 &kind,
836 Ok(Some(serde_json::json!({"accepted": true}))),
837 ),
838 )
839 .await;
840 let outcome_cancel = cancel_task.clone();
841 let execute_timing = RpcPhaseTiming::start(phase_diagnostics, "execute");
842 let result = if compact {
843 rpc_compact_cancellable(&lane_task, instructions, cancel_task).await
844 } else if resume {
845 rpc_resume_cancellable(&lane_task, cancel_task).await
846 } else {
847 rpc_prompt_cancellable(&lane_task, content, cancel_task).await
848 };
849 execute_timing.finish();
850 let flush_timing = RpcPhaseTiming::start(phase_diagnostics, "event_flush");
851 if let Some(flush) = flush_events {
852 let (done, drained) = tokio::sync::oneshot::channel();
853 if flush.send(done).is_ok() {
854 let _ = drained.await;
855 }
856 }
857 flush_timing.finish();
858 let (outcome, aborted, error) = match &result {
859 Ok(value) => {
860 let outcome = value
861 .get("outcome")
862 .and_then(|v| v.as_str())
863 .unwrap_or("completed")
864 .to_string();
865 let error = value
866 .get("error")
867 .and_then(|v| v.as_str())
868 .map(str::to_owned);
869 (outcome.clone(), outcome == "aborted", error)
870 }
871 Err(error) if outcome_cancel.is_cancelled() => {
872 ("aborted".to_string(), true, Some(error.clone()))
873 }
874 Err(ref error) => ("failed".to_string(), false, Some(error.clone())),
875 };
876 let resume_timing =
879 RpcPhaseTiming::start(phase_diagnostics, "pending_resume");
880 let pending_resume = lane_task.has_pending_resume().await;
881 resume_timing.finish();
882 let summary = result
883 .as_ref()
884 .ok()
885 .and_then(|value| value.get("summary"))
886 .cloned();
887 let settled = serde_json::json!({
888 "type": "agent_settled",
889 "runId": result.as_ref().ok().and_then(|value| value.get("runId")),
890 "outcome": outcome,
891 "aborted": aborted,
892 "error": error,
893 "summary": summary,
894 "pendingResume": pending_resume,
895 "reclaimed": result
899 .as_ref()
900 .ok()
901 .and_then(|value| value.get("reclaimed"))
902 .cloned(),
903 });
904 let enqueue_timing =
905 RpcPhaseTiming::start(phase_diagnostics, "settled_enqueue");
906 write_line(&out_task, &settled).await;
907 enqueue_timing.finish();
908 mailbox_task.cancel_all();
909 active_task.store(false, std::sync::atomic::Ordering::Release);
910 }),
911 run_cancel.clone(),
912 ));
913 continue;
914 }
915 _ => {}
916 }
917
918 let result = match kind.as_str() {
919 "abort" if run_active.load(std::sync::atomic::Ordering::Acquire) => {
920 run_cancel.cancel();
923 let result = lane.abort().await;
924 external_tools.abort();
925 abort_result_reply(result)
929 }
930 "abort" => {
931 let result = lane.abort().await;
932 external_tools.abort();
933 abort_result_reply(result)
934 }
935 "configure_external_tools" => {
936 rpc_configure_external_tools(
937 &lane,
938 harness,
939 run_active.load(std::sync::atomic::Ordering::Acquire),
940 &external_tools,
941 &mut external_names,
942 &msg,
943 )
944 .await
945 }
946 "external_tool_response" => external_tools
947 .respond(&msg)
948 .map(|accepted| serde_json::json!({"accepted": accepted})),
949 "configure_file_io" => {
950 rpc_configure_file_io(
951 &lane,
952 run_active.load(std::sync::atomic::Ordering::Acquire),
953 file_io_router.as_ref(),
954 &mut owner.file_io,
955 &stdout,
956 &msg,
957 )
958 .await
959 }
960 "file_io_response" => {
962 let mut reply = msg;
966 if let Some(object) = reply.as_object_mut() {
967 object.remove("id");
968 }
969 Ok(serde_json::json!({"accepted": owner.file_io.as_ref()
970 .is_some_and(|client| client.resolve_response(&reply))}))
971 }
972 "get_resource_state" => {
973 if msg.as_object().is_none_or(|o| {
974 o.keys()
975 .any(|k| !["type", "id", "version"].contains(&k.as_str()))
976 }) || msg.get("version").is_some_and(|v| v.as_u64() != Some(1))
977 {
978 Err("unsupported resource state request/schema (version 1 only)".into())
979 } else {
980 crate::resource_selection::get_resource_state(harness).await
981 }
982 }
983 "get_state" => Ok(rpc_get_state(&lane, harness).await),
984 "get_commands" => Ok(rpc_get_commands(oauth_snapshot.as_deref())),
985 "get_messages" => rpc_get_messages(&lane, harness, &msg).await,
986 "get_message_summaries" => crate::rpc::message_queries::summaries(harness, &msg).await,
987 "get_message_window" => crate::rpc::message_queries::window(harness, &msg).await,
988 "get_history" => rpc_get_history(&lane, harness, &msg).await,
989 "get_session_branch" => rpc_get_session_branch(harness, &msg).await,
990 "append_entry" => rpc_append_entry(harness, &msg).await,
991 "get_stats" => rpc_get_stats(&lane, harness).await,
992 "get_queue" => rpc_get_queue(&lane, harness, &msg).await,
993 "queue_message" => rpc_queue_message(&lane, &msg).await,
994 "cancel_queued" => rpc_cancel_queued(&lane, &msg).await,
995 "clear_queue" => rpc_clear_queue(&lane).await,
996 "navigate" => rpc_navigate(&lane, &msg).await,
997 "get_available_models" => Ok(rpc_get_available_models(&model_catalog)),
998 "get_auth_status" => Ok(rpc_get_auth_status(&model_catalog, oauth_snapshot.as_ref())),
999 "oauth" => rpc_oauth(&oauth_snapshot, &msg),
1000 "set_api_key" => rpc_set_api_key(&msg).await,
1001 "set_model" if run_active.load(std::sync::atomic::Ordering::Acquire) => {
1002 Err("cannot set model while a run or prompt validation is active".into())
1003 }
1004 "set_model" => rpc_set_model(&lane, &model_catalog, &msg).await,
1005 "get_available_thinking_levels" => Ok(rpc_get_available_thinking_levels(&lane).await),
1006 "set_thinking_level" => rpc_set_thinking_level(&lane, &msg).await,
1007 "get_available_tools" => Ok(rpc_get_available_tools(&lane, harness).await),
1008 "set_active_tools" => rpc_set_active_tools(&lane, &msg).await,
1009 "ping" => Ok(serde_json::json!({"pong": true})),
1010 other => Err(format!("unknown command type: {other}")),
1011 };
1012
1013 let response = rpc_protocol::response(Some(id.clone()), &kind, result.map(Some));
1014 write_line(&stdout, &response).await;
1015 }
1016
1017 ui_dialog.detach();
1021 if cancel_on_exit {
1022 run_cancel.cancel();
1023 let _ = lane.abort().await;
1024 }
1025 external_tools.shutdown();
1028 if let Some(client) = owner.file_io.take() {
1029 client.close();
1030 }
1031 for task in command_tasks {
1032 let _ = task.await;
1033 }
1034 let mut exit_code = 0;
1035 if let Some(mut task) = run_task {
1037 let result = tokio::select! {
1038 biased;
1039 result = &mut task => result,
1040 _ = writer_failed.cancelled() => {
1041 run_cancel.cancel();
1042 let _ = lane.abort().await;
1043 external_tools.shutdown();
1044 task.await
1045 }
1046 };
1047 if let Err(error) = result {
1048 eprintln!("[rpi] rpc prompt task failed: {error}");
1049 exit_code = 1;
1050 }
1051 }
1052 pump_stop.cancel();
1053 harness.session().set_entry_commit_publisher(None);
1054 drop(commit_tx);
1057 if let Err(error) = commit_task.await {
1058 eprintln!("[rpi] rpc commit event task failed: {error}");
1059 exit_code = 1;
1060 }
1061 let _ = dialog_task.await;
1062 if let Some(task) = event_task {
1063 let _ = task.await;
1064 }
1065 drop(stdout);
1067 match writer_task.await {
1068 Ok(Ok(())) => {}
1069 error => {
1070 eprintln!("[rpi] rpc output failed: {error:?}");
1071 exit_code = 1;
1072 }
1073 }
1074 owner.armed = false;
1075 exit_code
1076}
1077
1078type OwnedRpcRun = RetainedRpcTask<()>;
1084struct RetainedRpcTask<T: Send + 'static> {
1085 task: Option<tokio::task::JoinHandle<T>>,
1086 cancel: tokio_util::sync::CancellationToken,
1087}
1088impl<T: Send + 'static> RetainedRpcTask<T> {
1089 fn new(task: tokio::task::JoinHandle<T>, cancel: tokio_util::sync::CancellationToken) -> Self {
1090 Self {
1091 task: Some(task),
1092 cancel,
1093 }
1094 }
1095 fn is_finished(&self) -> bool {
1096 self.task.as_ref().is_none_or(|task| task.is_finished())
1097 }
1098}
1099impl<T: Send + 'static> std::future::Future for RetainedRpcTask<T> {
1100 type Output = Result<T, tokio::task::JoinError>;
1101 fn poll(
1102 mut self: std::pin::Pin<&mut Self>,
1103 cx: &mut std::task::Context<'_>,
1104 ) -> std::task::Poll<Self::Output> {
1105 let result = std::future::Future::poll(
1106 std::pin::Pin::new(self.task.as_mut().expect("RPC task already joined")),
1107 cx,
1108 );
1109 if result.is_ready() {
1110 self.task.take();
1111 }
1112 result
1113 }
1114}
1115impl<T: Send + 'static> Drop for RetainedRpcTask<T> {
1116 fn drop(&mut self) {
1117 if let Some(task) = self.task.take() {
1118 if !task.is_finished() {
1120 self.cancel.cancel();
1121 }
1122 tokio::spawn(async move {
1123 if let Err(error) = task.await {
1124 eprintln!("[rpi] cancelled RPC task cleanup failed: {error}");
1125 }
1126 });
1127 }
1128 }
1129}
1130
1131struct RpcOwner<'a> {
1134 harness: &'a AgentHarness,
1135 file_io: Option<crate::rpc::file_io::FileIOClient>,
1136 backend: Arc<crate::rpc::external_tools::ExternalTools>,
1137 run_cancel: tokio_util::sync::CancellationToken,
1138 pump_stop: tokio_util::sync::CancellationToken,
1139 failed: tokio_util::sync::CancellationToken,
1140 mailbox: rpi_extensions::UiDialogMailbox,
1141 armed: bool,
1142}
1143impl Drop for RpcOwner<'_> {
1144 fn drop(&mut self) {
1145 if self.armed {
1146 self.run_cancel.cancel();
1147 if let Some(client) = self.file_io.take() {
1148 client.close();
1149 }
1150 self.backend.shutdown();
1151 self.failed.cancel();
1152 self.pump_stop.cancel();
1153 self.mailbox.detach();
1154 self.harness.session().set_entry_commit_publisher(None);
1155 }
1156 }
1157}
1158
1159fn rpc_get_commands(snapshot: Option<&rpi_extensions::RegistrySnapshot>) -> serde_json::Value {
1160 let commands: Vec<_> = snapshot
1161 .into_iter()
1162 .flat_map(|s| s.commands())
1163 .map(
1164 |command| serde_json::json!({"name": command.name, "description": command.description}),
1165 )
1166 .collect();
1167 serde_json::json!({"commands": commands})
1168}
1169
1170fn rpc_execute_command(
1171 snapshot: Option<&rpi_extensions::RegistrySnapshot>,
1172 request: &serde_json::Value,
1173) -> Result<serde_json::Value, String> {
1174 let name = request
1175 .get("command")
1176 .and_then(|v| v.as_str())
1177 .filter(|name| !name.trim_start_matches('/').is_empty())
1178 .ok_or("command must be a non-empty string")?;
1179 let args = match request.get("args") {
1180 None => "",
1181 Some(value) => value.as_str().ok_or("args must be a string")?,
1182 };
1183 let command = snapshot
1184 .into_iter()
1185 .flat_map(|s| s.commands())
1186 .find(|c| c.name.trim_start_matches('/') == name.trim_start_matches('/'))
1187 .ok_or_else(|| format!("unknown command: {name}"))?;
1188 let input = serde_json::json!({"command": name, "args": args}).to_string();
1189 std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
1190 let mut out = rpi_plugin_sdk::StbString::empty();
1191 let rc = (command.handler)(
1192 rpi_plugin_sdk::StbStringRef::from_str(&input),
1193 &mut out,
1194 command.user_data,
1195 );
1196 let out = unsafe {
1199 rpi_plugin_sdk::callbacks::OwnedString::from_raw(
1200 out,
1201 Some(rpi_extensions::host_free_string),
1202 )
1203 };
1204 if rc != 0 {
1205 return Err(format!("command failed: {name} (code {rc})"));
1206 }
1207 serde_json::from_str(&out.to_string_lossy())
1208 .map_err(|_| format!("command returned invalid JSON: {name}"))
1209 }))
1210 .unwrap_or_else(|_| Err(format!("command panicked: {name}")))
1211}
1212
1213async fn drain_rpc_events(
1214 out: &RpcWriter,
1215 events: &mut tokio::sync::broadcast::Receiver<AgentEvent>,
1216) {
1217 loop {
1218 match events.try_recv() {
1219 Ok(event) => write_agent_event(out, &event).await,
1220 Err(tokio::sync::broadcast::error::TryRecvError::Lagged(_)) => continue,
1221 Err(_) => break,
1222 }
1223 }
1224}
1225
1226async fn rpc_validate_prompt(
1229 lane: &Arc<dyn AgentLane>,
1230 command: serde_json::Value,
1231 cancel: &tokio_util::sync::CancellationToken,
1232) -> Result<crate::rpc::prompt_content::ParsedPrompt, String> {
1233 use crate::rpc::prompt_content::parse_prompt_content;
1234 if cancel.is_cancelled() {
1235 return Err("Operation aborted during prompt validation".into());
1236 }
1237 let model = lane.get_model().await.map_err(|e| e.to_string())?;
1241 let mut worker = tokio::task::spawn_blocking(move || parse_prompt_content(&command));
1242 let result = tokio::select! {
1243 biased;
1244 _ = cancel.cancelled() => {
1245 let _ = worker.await;
1248 return Err("Operation aborted during prompt validation".into());
1249 }
1250 result = &mut worker => result.map_err(|_| "prompt validation worker failed".to_string())?,
1251 }?;
1252 if cancel.is_cancelled() {
1253 return Err("Operation aborted during prompt validation".into());
1254 }
1255 if result.has_images() {
1256 let current = lane.get_model().await.map_err(|e| e.to_string())?;
1257 if current.id != model.id || current.provider != model.provider || current.api != model.api
1258 {
1259 return Err("model changed during image prompt validation; retry while idle".into());
1260 }
1261 if !model.input.contains(&rpi_ai::types::InputModality::Image)
1262 || !current.input.contains(&rpi_ai::types::InputModality::Image)
1263 {
1264 return Err("selected model does not support image input".into());
1265 }
1266 }
1267 Ok(result)
1268}
1269
1270async fn rpc_configure_file_io(
1271 lane: &Arc<dyn AgentLane>,
1272 run_active: bool,
1273 router: Option<&Arc<rpi_tools::file_io::FileIORouter>>,
1274 owner: &mut Option<crate::rpc::file_io::FileIOClient>,
1275 writer: &RpcWriter,
1276 command: &serde_json::Value,
1277) -> Result<serde_json::Value, String> {
1278 use rpi_tools::file_io::FileIOCapabilities;
1279 if run_active
1280 || lane
1281 .snapshot()
1282 .await
1283 .map_err(|e| e.to_string())?
1284 .active
1285 .is_some()
1286 {
1287 return Err("cannot configure file I/O while a run or prompt validation is active".into());
1288 }
1289 let read_text = command
1290 .get("readText")
1291 .and_then(serde_json::Value::as_bool)
1292 .ok_or("readText must be a boolean")?;
1293 let write_text = command
1294 .get("writeText")
1295 .and_then(serde_json::Value::as_bool)
1296 .ok_or("writeText must be a boolean")?;
1297 let router = router.ok_or("this embedded harness has no attached built-in file I/O router")?;
1298 if let Some(old) = owner.take() {
1301 old.close();
1302 }
1303 let next =
1304 (read_text || write_text).then(|| crate::rpc::file_io::FileIOClient::new(writer.clone()));
1305 router
1306 .configure(
1307 FileIOCapabilities {
1308 read_text,
1309 write_text,
1310 },
1311 next.as_ref().map(|c| c.backend()),
1312 )
1313 .await
1314 .map_err(|e| e.to_string())?;
1315 *owner = next;
1316 Ok(serde_json::json!({"readText": read_text, "writeText": write_text}))
1317}
1318
1319async fn rpc_prompt_cancellable(
1322 lane: &Arc<dyn AgentLane>,
1323 content: crate::rpc::prompt_content::ParsedPrompt,
1324 cancel: tokio_util::sync::CancellationToken,
1325) -> Result<serde_json::Value, String> {
1326 if cancel.is_cancelled() {
1327 return Err("Operation aborted".into());
1328 }
1329 let prompt = rpc_prompt_content(lane, content);
1330 tokio::pin!(prompt);
1331 tokio::select! {
1332 biased;
1333 _ = cancel.cancelled() => {
1334 match lane.abort().await {
1338 Ok(reclaimed) => prompt.await.map(|mut value| {
1339 if let Some(object) = value.as_object_mut() {
1340 object.insert(
1341 "reclaimed".to_string(),
1342 reclaimed_payload(&reclaimed),
1343 );
1344 }
1345 value
1346 }),
1347 Err(_) => Err("Operation aborted".to_string()),
1348 }
1349 }
1350 result = &mut prompt => result,
1351 }
1352}
1353
1354fn reclaimed_payload(result: &rpi_harness::agent_harness::AbortResult) -> serde_json::Value {
1356 serde_json::json!({
1357 "steer": crate::extensions_actions::queued_message_payloads(&result.steer),
1358 "followUp": crate::extensions_actions::queued_message_payloads(&result.follow_up),
1359 })
1360}
1361
1362async fn rpc_compact_cancellable(
1368 lane: &Arc<dyn AgentLane>,
1369 instructions: Option<String>,
1370 cancel: tokio_util::sync::CancellationToken,
1371) -> Result<serde_json::Value, String> {
1372 let compact = async {
1373 match lane.compact(instructions.as_deref()).await {
1374 Ok(result) => Ok(compaction_json(&result)),
1375 Err(error) => Err(error.to_string()),
1376 }
1377 };
1378 tokio::pin!(compact);
1379 tokio::select! {
1380 biased;
1381 _ = cancel.cancelled() => {
1382 let _ = lane.abort().await;
1383 compact.await
1384 }
1385 result = &mut compact => result,
1386 }
1387}
1388
1389fn compaction_json(result: &CompactionResult) -> serde_json::Value {
1392 let mut value = match &result.outcome {
1393 CompactionOutcome::Completed { leaf_id, entry } => {
1394 let (tokens_before, retained) = match entry {
1395 rpi_harness::session::types::Entry::Compaction(compaction) => {
1396 (compaction.tokens_before, compaction.retained_tail.len())
1397 }
1398 _ => (0, 0),
1399 };
1400 serde_json::json!({
1401 "outcome": "compacted",
1402 "leafId": leaf_id,
1403 "tokensBefore": tokens_before,
1404 "retainedMessages": retained,
1405 "summary": format!(
1406 "Conversation compacted · {} tokens before, {} messages kept.",
1407 tokens_before, retained
1408 ),
1409 })
1410 }
1411 CompactionOutcome::Declined { leaf_id } => serde_json::json!({
1412 "outcome": "declined",
1413 "leafId": leaf_id,
1414 "summary": "Nothing to compact yet.",
1415 }),
1416 CompactionOutcome::Aborted { leaf_id } => serde_json::json!({
1417 "outcome": "aborted",
1418 "leafId": leaf_id,
1419 "summary": "Compaction aborted.",
1420 }),
1421 CompactionOutcome::Failed { leaf_id, error } => serde_json::json!({
1422 "outcome": "failed",
1423 "leafId": leaf_id,
1424 "error": error.message,
1425 "summary": format!("Compaction failed: {}", error.message),
1426 }),
1427 };
1428 value["runId"] = serde_json::json!(result.run_id);
1429 value
1430}
1431
1432async fn rpc_resume_cancellable(
1437 lane: &Arc<dyn AgentLane>,
1438 cancel: tokio_util::sync::CancellationToken,
1439) -> Result<serde_json::Value, String> {
1440 let resume = async {
1441 match lane.resume_pending().await {
1442 Ok(Some(result)) => {
1443 let (outcome, final_text, error) = rpc_outcome_str(&result.outcome);
1444 Ok(serde_json::json!({
1445 "runId": result.run_id,
1446 "outcome": outcome,
1447 "finalText": final_text,
1448 "error": error,
1449 "resumed": true,
1450 }))
1451 }
1452 Ok(None) => Ok(serde_json::json!({"outcome": "skipped", "resumed": false})),
1453 Err(error) => Err(error.to_string()),
1454 }
1455 };
1456 tokio::pin!(resume);
1457 tokio::select! {
1458 biased;
1459 _ = cancel.cancelled() => {
1460 let _ = lane.abort().await;
1461 resume.await
1462 }
1463 result = &mut resume => result,
1464 }
1465}
1466
1467fn rpc_outcome_str(outcome: &HarnessRunOutcome) -> (&'static str, Option<String>, Option<String>) {
1468 match outcome {
1469 HarnessRunOutcome::Completed { final_message, .. } => {
1470 ("completed", Some(assistant_text(final_message)), None)
1471 }
1472 HarnessRunOutcome::Aborted { final_message, .. } => {
1473 ("aborted", Some(assistant_text(final_message)), None)
1474 }
1475 HarnessRunOutcome::Failed {
1478 final_message,
1479 error,
1480 ..
1481 } => (
1482 "failed",
1483 final_message.as_ref().map(assistant_text),
1484 Some(error.message.clone()),
1485 ),
1486 HarnessRunOutcome::Suspended { .. } => ("suspended", None, None),
1487 }
1488}
1489
1490async fn rpc_prompt_content(
1491 lane: &Arc<dyn AgentLane>,
1492 content: crate::rpc::prompt_content::ParsedPrompt,
1493) -> Result<serde_json::Value, String> {
1494 use crate::rpc::prompt_content::ParsedPrompt;
1495 let result = match content {
1496 ParsedPrompt::Legacy(text) => lane.prompt_text(&text, Vec::new()).await,
1497 ParsedPrompt::Blocks(blocks) => {
1498 lane.prompt_message(rpi_agent::AgentMessage::User(UserMessage::new(
1499 rpi_ai::types::UserContent::Blocks(blocks),
1500 rpc_now_ms(),
1501 )))
1502 .await
1503 }
1504 }
1505 .map_err(|e| e.to_string())?;
1506 let (outcome, final_text, error) = rpc_outcome_str(&result.outcome);
1507 Ok(
1508 serde_json::json!({"runId":result.run_id, "outcome": outcome, "finalText": final_text, "error": error}),
1509 )
1510}
1511
1512async fn rpc_get_state(lane: &Arc<dyn AgentLane>, harness: &AgentHarness) -> serde_json::Value {
1513 let model = lane.get_model().await.ok();
1514 let thinking = lane
1515 .get_thinking_level()
1516 .await
1517 .ok()
1518 .map(|t| serde_json::to_value(t).unwrap_or(serde_json::Value::Null));
1519 let active_tools = if harness.has_exact_resources() {
1520 harness
1521 .execution_context()
1522 .ok()
1523 .and_then(|v| v.get("tools").cloned())
1524 .unwrap_or_else(|| serde_json::json!([]))
1525 } else {
1526 serde_json::json!(lane.get_active_tools().await.unwrap_or_default())
1527 };
1528 let leaf_id = lane.get_leaf_id().await.ok().flatten();
1529 let context = rpc_context_usage(lane, harness).await;
1530 let active_operation = lane.active_operation().await.ok().flatten().map(|active| {
1531 serde_json::json!({
1532 "runId": active.run_id,
1533 "lane": active.lane,
1534 "kind": format!("{:?}", active.kind).to_lowercase(),
1535 })
1536 });
1537 let queued = lane.queued_messages().await.unwrap_or_default();
1538 let stats = rpc_get_stats(lane, harness).await.ok();
1539 serde_json::json!({
1540 "modelId": model.as_ref().map(|model| model.id.clone()),
1541 "model": model.as_ref().map(|model| model.id.clone()),
1542 "provider": model.as_ref().map(|model| model.provider.clone()),
1543 "thinkingLevel": thinking,
1544 "activeTools": active_tools,
1545 "leafId": leaf_id,
1546 "context": context,
1547 "sessionId": harness.session().storage().metadata().id,
1548 "activeOperation": active_operation,
1549 "queue": {
1550 "steering": queued
1558 .steering
1559 .iter()
1560 .map(|m| m.text.as_str())
1561 .collect::<Vec<_>>(),
1562 "followUp": queued.follow_up.iter().map(|m| m.text.as_str()).collect::<Vec<_>>(),
1563 "steeringCount": queued.steering.len(),
1564 "followUpCount": queued.follow_up.len(),
1565 },
1566 "stats": stats,
1567 })
1568}
1569
1570fn rpc_now_ms() -> i64 {
1571 std::time::SystemTime::now()
1572 .duration_since(std::time::UNIX_EPOCH)
1573 .map(|duration| duration.as_millis() as i64)
1574 .unwrap_or_default()
1575}
1576
1577fn rpc_message_text(value: &serde_json::Value) -> String {
1578 let Some(content) = value.get("content") else {
1579 return String::new();
1580 };
1581 if let Some(text) = content.as_str() {
1582 return text.to_string();
1583 }
1584 content
1585 .as_array()
1586 .into_iter()
1587 .flatten()
1588 .filter_map(|block| block.get("text").and_then(|text| text.as_str()))
1589 .collect::<Vec<_>>()
1590 .join("\n")
1591}
1592
1593async fn rpc_get_history(
1594 lane: &Arc<dyn AgentLane>,
1595 harness: &AgentHarness,
1596 msg: &serde_json::Value,
1597) -> Result<serde_json::Value, String> {
1598 let limit = msg
1599 .get("limit")
1600 .and_then(serde_json::Value::as_u64)
1601 .unwrap_or(100)
1602 .clamp(1, 500) as usize;
1603 let after_seq = msg.get("afterSeq").and_then(serde_json::Value::as_u64);
1604 let before_seq = msg.get("beforeSeq").and_then(serde_json::Value::as_u64);
1605 let page = harness
1606 .session()
1607 .record_page(&rpi_harness::session::types::RecordPageQuery {
1608 lane: lane.name().to_owned(),
1609 after_seq,
1610 before_seq,
1611 limit,
1612 from_end: before_seq.is_some() || after_seq.is_none(),
1613 })
1614 .await
1615 .map_err(|error| error.to_string())?;
1616 let records = page.records;
1617 let total = page.total;
1618 let has_more = total > records.len();
1619 let first_seq = records.first().map(LaneRecord::seq);
1620 let last_seq = records.last().map(LaneRecord::seq);
1621 let open_operations = harness
1622 .session()
1623 .find_open_operations(lane.name(), Some(50))
1624 .await
1625 .map_err(|error| error.to_string())?;
1626 Ok(serde_json::json!({
1627 "records": records.iter().map(|record| serde_json::to_value(record).unwrap_or(serde_json::Value::Null)).collect::<Vec<_>>(),
1628 "total": total,
1629 "limit": limit,
1630 "afterSeq": after_seq,
1631 "beforeSeq": before_seq,
1632 "firstSeq": first_seq,
1633 "lastSeq": last_seq,
1634 "hasMore": has_more,
1635 "leafId": lane.get_leaf_id().await.ok().flatten(),
1636 "openOperations": open_operations.iter().map(|record| serde_json::to_value(record).unwrap_or(serde_json::Value::Null)).collect::<Vec<_>>(),
1637 }))
1638}
1639
1640async fn rpc_get_session_branch(
1641 harness: &AgentHarness,
1642 msg: &serde_json::Value,
1643) -> Result<serde_json::Value, String> {
1644 let mut args = msg.clone();
1647 if let Some(args) = args.as_object_mut() {
1648 args.remove("type");
1649 args.remove("id");
1650 }
1651 let request = branch_entry_query(&args)?;
1652 let (entries, session_id, leaf_id) = branch_snapshot(harness, &request).await?;
1653 let entries = entries
1654 .iter()
1655 .map(|entry| serde_json::to_value(entry).map_err(|error| error.to_string()))
1656 .collect::<Result<Vec<_>, _>>()?;
1657 Ok(serde_json::json!({
1658 "entries": entries,
1659 "count": entries.len(),
1660 "filtered": true,
1661 "sessionId": session_id,
1662 "leafId": leaf_id,
1663 "scopedAppend": true,
1664 }))
1665}
1666
1667#[derive(Debug, Clone, PartialEq, Eq)]
1668pub(crate) struct BranchEntryQuery {
1669 pub custom_types: Vec<String>,
1670 pub order: EntryOrder,
1671 pub limit: Option<usize>,
1672 pub after_seq: Option<u64>,
1673}
1674
1675impl BranchEntryQuery {
1676 fn entry_queries(&self) -> Vec<EntryQuery> {
1677 let per_query_limit = self.limit;
1680 self.custom_types
1681 .iter()
1682 .map(|custom_type| EntryQuery {
1683 entry_type: Some("custom"),
1684 custom_type: Some(custom_type.clone()),
1685 order: Some(self.order),
1686 limit: per_query_limit,
1687 cursor: self.after_seq.map(|after_seq| EntryCursor { after_seq }),
1688 ..Default::default()
1689 })
1690 .collect()
1691 }
1692}
1693
1694pub(crate) async fn branch_snapshot(
1695 harness: &AgentHarness,
1696 request: &BranchEntryQuery,
1697) -> Result<(Vec<Entry>, String, Option<String>), String> {
1698 let storage = harness.session().storage();
1702 let leaf_id = storage
1703 .get_lanes()
1704 .await
1705 .map_err(|error| error.to_string())?
1706 .into_iter()
1707 .find(|pointer| pointer.lane == "main")
1708 .ok_or_else(|| "Lane not found: main".to_string())?
1709 .leaf_id;
1710 let session_id = storage.metadata().id;
1711 let Some(start) = leaf_id.clone() else {
1712 return Ok((Vec::new(), session_id, None));
1713 };
1714 let bounds = BranchBounds {
1715 start: Some(start.clone()),
1716 ..Default::default()
1717 };
1718 let mut entries = Vec::new();
1719 for query in request.entry_queries() {
1720 entries.extend(
1721 storage
1722 .find_entries_on_branch(&query, &bounds, &start)
1723 .await
1724 .map_err(|error| error.to_string())?,
1725 );
1726 }
1727 if request.custom_types.len() > 1 {
1728 let mut seen = HashSet::with_capacity(entries.len());
1729 entries.retain(|entry| seen.insert(entry.id().to_owned()));
1730 match request.order {
1731 EntryOrder::OldestFirst => entries.sort_by_key(Entry::seq),
1732 EntryOrder::NewestFirst => entries.sort_by_key(|entry| Reverse(entry.seq())),
1733 }
1734 if let Some(limit) = request.limit {
1735 entries.truncate(limit);
1736 }
1737 }
1738 Ok((entries, session_id, leaf_id))
1739}
1740
1741pub(crate) fn branch_entry_query(msg: &serde_json::Value) -> Result<BranchEntryQuery, String> {
1744 let args = msg
1745 .as_object()
1746 .ok_or("get_session_branch: query must be an object")?;
1747 for key in args.keys() {
1748 if !matches!(key.as_str(), "customTypes" | "afterSeq" | "limit" | "order") {
1749 return Err(format!(
1750 "get_session_branch: unknown query argument `{key}`"
1751 ));
1752 }
1753 }
1754 let values = args
1755 .get("customTypes")
1756 .and_then(serde_json::Value::as_array)
1757 .filter(|values| !values.is_empty())
1758 .ok_or(
1759 "get_session_branch: `customTypes` is required and must be a non-empty string array",
1760 )?;
1761 let mut custom_types = Vec::with_capacity(values.len());
1762 for value in values {
1763 let custom_type = value
1764 .as_str()
1765 .filter(|value| !value.trim().is_empty())
1766 .ok_or("get_session_branch: every `customTypes` item must be a non-empty string")?;
1767 if !custom_types.iter().any(|existing| existing == custom_type) {
1768 custom_types.push(custom_type.to_owned());
1769 }
1770 }
1771 let limit = match msg.get("limit") {
1772 None | Some(serde_json::Value::Null) => None,
1773 Some(value) => {
1774 let limit = value
1775 .as_u64()
1776 .filter(|limit| *limit > 0)
1777 .ok_or("get_session_branch: `limit` must be a positive integer")?;
1778 Some(usize::try_from(limit).map_err(|_| "get_session_branch: `limit` is too large")?)
1779 }
1780 };
1781 let order = match msg.get("order") {
1782 None => EntryOrder::OldestFirst,
1783 Some(serde_json::Value::String(value)) if value == "oldestFirst" => EntryOrder::OldestFirst,
1784 Some(serde_json::Value::String(value)) if value == "newestFirst" => EntryOrder::NewestFirst,
1785 Some(_) => {
1786 return Err("get_session_branch: `order` must be oldestFirst or newestFirst".into())
1787 }
1788 };
1789 let after_seq = match msg.get("afterSeq") {
1790 None | Some(serde_json::Value::Null) => None,
1791 Some(value) => Some(
1792 value
1793 .as_u64()
1794 .ok_or("get_session_branch: `afterSeq` must be a non-negative integer")?,
1795 ),
1796 };
1797 Ok(BranchEntryQuery {
1798 custom_types,
1799 order,
1800 limit,
1801 after_seq,
1802 })
1803}
1804
1805async fn rpc_append_entry(
1806 harness: &AgentHarness,
1807 msg: &serde_json::Value,
1808) -> Result<serde_json::Value, String> {
1809 let custom_type = msg
1810 .get("customType")
1811 .and_then(serde_json::Value::as_str)
1812 .filter(|value| !value.trim().is_empty())
1813 .ok_or_else(|| "append_entry requires a non-empty string `customType`".to_string())?;
1814 crate::resource_selection::validate_public_entry(custom_type, msg.get("data"))?;
1815 let entry_id = if let Some(expected) = msg.get("expected") {
1816 let guard = serde_json::from_value(expected.clone())
1817 .map_err(|error| format!("invalid expected guard: {error}"))?;
1818 harness
1819 .append_scoped_entry(
1820 &guard,
1821 rpi_harness::agent_harness::ScopedEntry::Custom {
1822 custom_type: custom_type.to_owned(),
1823 data: msg.get("data").cloned(),
1824 },
1825 )
1826 .await
1827 } else {
1828 harness
1829 .append_custom_entry(custom_type, msg.get("data").cloned())
1830 .await
1831 }
1832 .map_err(|error| error.to_string())?;
1833 let mut response = serde_json::json!({"entryId": entry_id, "customType": custom_type});
1834 if msg.get("expected").is_some() {
1835 response["scoped"] = serde_json::json!(true);
1836 }
1837 Ok(response)
1838}
1839
1840async fn rpc_get_stats(
1841 lane: &Arc<dyn AgentLane>,
1842 harness: &AgentHarness,
1843) -> Result<serde_json::Value, String> {
1844 let stats = lane.get_stats().await.map_err(|error| error.to_string())?;
1845 let mut records = Vec::new();
1847 for record_type in [
1848 "operation_started",
1849 "operation_finished",
1850 "retry_pending",
1851 "usage",
1852 ] {
1853 records.extend(
1854 harness
1855 .session()
1856 .find_records(&RecordQuery {
1857 lane: Some(lane.name().to_string()),
1858 record_type: Some(record_type),
1859 order: Some(EntryOrder::OldestFirst),
1860 ..Default::default()
1861 })
1862 .await
1863 .map_err(|error| error.to_string())?,
1864 );
1865 }
1866 records.sort_unstable_by_key(LaneRecord::seq);
1867 let mut operation_counts = std::collections::BTreeMap::<String, u64>::new();
1868 for record in &records {
1869 let kind = match record {
1870 LaneRecord::OperationStarted(started) => {
1871 let key = started.intent.kind();
1872 *operation_counts.entry(key.to_string()).or_default() += 1;
1873 if matches!(
1874 &started.intent,
1875 rpi_harness::session::types::OperationIntent::Run {
1876 resume_data: Some(_),
1877 ..
1878 }
1879 ) {
1880 *operation_counts.entry("resume".to_string()).or_default() += 1;
1881 }
1882 continue;
1883 }
1884 LaneRecord::OperationFinished(_) => "operation_finished",
1885 LaneRecord::RetryPending(_) => "retry_pending",
1886 LaneRecord::Usage(_) => "usage",
1887 _ => continue,
1888 };
1889 *operation_counts.entry(kind.to_string()).or_default() += 1;
1890 }
1891 let count = |name: &str| operation_counts.get(name).copied().unwrap_or_default();
1892 let model = lane.get_model().await.ok();
1893 let input = stats.cached_tokens + stats.uncached_tokens;
1894 Ok(serde_json::json!({
1895 "messageCount": stats.message_count,
1896 "operationCount": count("operation_finished"),
1897 "retryCount": count("retry_pending"),
1898 "compactionCount": count("compaction"),
1899 "resumeCount": count("resume"),
1900 "operationCounts": operation_counts,
1901 "provider": model.as_ref().map(|value| value.provider.clone()),
1902 "modelId": model.as_ref().map(|value| value.id.clone()),
1903 "cachedTokens": stats.cached_tokens,
1904 "uncachedTokens": stats.uncached_tokens,
1905 "totalTokens": stats.total_tokens,
1906 "costTotal": stats.cost_total,
1907 "cacheHitRatio": if input > 0 { Some(stats.cached_tokens as f64 / input as f64) } else { None::<f64> },
1908 }))
1909}
1910
1911async fn rpc_get_queue(
1912 lane: &Arc<dyn AgentLane>,
1913 harness: &AgentHarness,
1914 msg: &serde_json::Value,
1915) -> Result<serde_json::Value, String> {
1916 let include_completed = msg
1917 .get("includeCompleted")
1918 .and_then(serde_json::Value::as_bool)
1919 .unwrap_or(false);
1920 let mut records = Vec::new();
1922 for record_type in ["queue_enqueued", "queue_cancelled"] {
1923 records.extend(
1924 harness
1925 .session()
1926 .find_records(&RecordQuery {
1927 lane: Some(lane.name().to_string()),
1928 record_type: Some(record_type),
1929 order: Some(EntryOrder::OldestFirst),
1930 ..Default::default()
1931 })
1932 .await
1933 .map_err(|error| error.to_string())?,
1934 );
1935 }
1936 records.sort_unstable_by_key(LaneRecord::seq);
1937 let mut cancelled = std::collections::HashMap::<String, u64>::new();
1938 for record in &records {
1939 if let LaneRecord::QueueCancelled(record) = record {
1940 cancelled.insert(record.entry_id.clone(), record.base.seq);
1941 }
1942 }
1943 let mut items = Vec::new();
1944 for record in records {
1945 let LaneRecord::QueueEnqueued(record) = record else {
1946 continue;
1947 };
1948 let Some(entry_id) = record.target.get("id").and_then(|value| value.as_str()) else {
1949 continue;
1950 };
1951 let (status, completed_seq) = if let Some(seq) = cancelled.get(entry_id) {
1952 ("cancelled", Some(*seq))
1953 } else if let Some(entry) = harness.get_entry(entry_id).await.ok().flatten() {
1954 ("consumed", Some(entry.base().seq))
1955 } else {
1956 ("queued", None)
1957 };
1958 if status != "queued" && !include_completed {
1959 continue;
1960 }
1961 let message = record
1962 .target
1963 .get("message")
1964 .cloned()
1965 .unwrap_or(serde_json::Value::Null);
1966 items.push(serde_json::json!({
1967 "entryId": entry_id,
1968 "queue": record.queue,
1969 "status": status,
1970 "text": rpc_message_text(&message),
1971 "message": message,
1972 "runId": record.run_id,
1973 "seq": record.base.seq,
1974 "completedSeq": completed_seq,
1975 }));
1976 }
1977 Ok(
1978 serde_json::json!({"items": items, "count": items.len(), "includeCompleted": include_completed}),
1979 )
1980}
1981
1982async fn rpc_queue_message(
1983 lane: &Arc<dyn AgentLane>,
1984 msg: &serde_json::Value,
1985) -> Result<serde_json::Value, String> {
1986 let text = msg
1987 .get("message")
1988 .and_then(serde_json::Value::as_str)
1989 .filter(|text| !text.trim().is_empty())
1990 .ok_or_else(|| "queue_message requires a non-empty string `message`".to_string())?;
1991 let message = rpi_agent::message::AgentMessage::User(UserMessage::new(text, rpc_now_ms()));
1992 let queue = msg
1993 .get("deliverAs")
1994 .and_then(serde_json::Value::as_str)
1995 .unwrap_or("followUp");
1996 let result = match queue {
1997 "steer" => lane.steer(message).await,
1998 "followUp" | "follow_up" => lane.follow_up(message).await,
1999 "nextTurn" | "next_run" => lane.next_run(message).await,
2000 other => return Err(format!("unknown queue kind: {other}")),
2001 }
2002 .map_err(|error| error.to_string())?;
2003 Ok(serde_json::json!({
2004 "accepted": true,
2005 "entryId": result.entry_id,
2006 "queue": queue,
2007 }))
2008}
2009
2010async fn rpc_cancel_queued(
2011 lane: &Arc<dyn AgentLane>,
2012 msg: &serde_json::Value,
2013) -> Result<serde_json::Value, String> {
2014 let entry_id = msg
2015 .get("entryId")
2016 .and_then(serde_json::Value::as_str)
2017 .filter(|value| !value.is_empty())
2018 .ok_or_else(|| "cancel_queued requires `entryId`".to_string())?;
2019 let result = lane
2020 .cancel_queued(entry_id)
2021 .await
2022 .map_err(|error| error.to_string())?;
2023 Ok(serde_json::json!({
2024 "entryId": entry_id,
2025 "outcome": format!("{:?}", result.outcome).to_lowercase(),
2026 }))
2027}
2028
2029fn abort_response(result: &rpi_harness::agent_harness::AbortResult) -> serde_json::Value {
2032 serde_json::json!({
2033 "aborted": true,
2034 "steer": crate::extensions_actions::queued_message_payloads(&result.steer),
2035 "followUp": crate::extensions_actions::queued_message_payloads(&result.follow_up),
2036 })
2037}
2038
2039fn abort_result_reply(
2045 result: Result<rpi_harness::agent_harness::AbortResult, rpi_harness::result::HarnessError>,
2046) -> Result<serde_json::Value, String> {
2047 match result {
2048 Ok(result) => Ok(abort_response(&result)),
2049 Err(rpi_harness::result::HarnessError::NoActiveRun { .. }) => {
2050 Ok(abort_response(&rpi_harness::agent_harness::AbortResult {
2051 run_id: String::new(),
2052 steer: Vec::new(),
2053 follow_up: Vec::new(),
2054 }))
2055 }
2056 Err(error) => Err(error.to_string()),
2057 }
2058}
2059
2060async fn rpc_clear_queue(lane: &Arc<dyn AgentLane>) -> Result<serde_json::Value, String> {
2061 let result = lane
2062 .clear_queue()
2063 .await
2064 .map_err(|error| error.to_string())?;
2065 Ok(serde_json::json!({
2066 "cleared": result.steering.len() + result.follow_up.len(),
2067 "steering": result
2070 .steering
2071 .iter()
2072 .map(|m| m.text.as_str())
2073 .collect::<Vec<_>>(),
2074 "followUp": result.follow_up.iter().map(|m| m.text.as_str()).collect::<Vec<_>>(),
2075 }))
2076}
2077
2078async fn rpc_navigate(
2079 lane: &Arc<dyn AgentLane>,
2080 msg: &serde_json::Value,
2081) -> Result<serde_json::Value, String> {
2082 let target_id = msg
2083 .get("targetId")
2084 .and_then(serde_json::Value::as_str)
2085 .filter(|value| !value.is_empty());
2086 if lane
2087 .snapshot()
2088 .await
2089 .map_err(|error| error.to_string())?
2090 .active
2091 .is_some()
2092 {
2093 return Err("cannot navigate while an operation is active".to_string());
2094 }
2095 let result = lane
2096 .navigate_tree(target_id, false, None, None)
2097 .await
2098 .map_err(|error| error.to_string())?;
2099 let (outcome, leaf_id, error) = match result.outcome {
2100 rpi_harness::agent_harness::NavigationOutcome::Completed { new_leaf_id, .. } => {
2101 ("completed", new_leaf_id, None)
2102 }
2103 rpi_harness::agent_harness::NavigationOutcome::Declined { leaf_id } => {
2104 ("declined", leaf_id, None)
2105 }
2106 rpi_harness::agent_harness::NavigationOutcome::Aborted { leaf_id } => {
2107 ("aborted", leaf_id, None)
2108 }
2109 rpi_harness::agent_harness::NavigationOutcome::Failed { leaf_id, error } => {
2110 ("failed", leaf_id, Some(error.message))
2111 }
2112 };
2113 Ok(serde_json::json!({
2114 "runId": result.run_id,
2115 "outcome": outcome,
2116 "leafId": leaf_id,
2117 "error": error,
2118 }))
2119}
2120
2121async fn rpc_context_usage(lane: &Arc<dyn AgentLane>, harness: &AgentHarness) -> serde_json::Value {
2127 use rpi_harness::compaction::{build_session_context, estimate_context_tokens};
2128 use rpi_harness::session::types::{BranchBounds, EntryOrder, EntryQuery};
2129
2130 let estimate = match harness.context_usage().await {
2131 Ok(Some(estimate)) => estimate,
2132 Ok(None) => {
2133 let entries = harness
2135 .find_entries_on_branch(
2136 &EntryQuery {
2137 order: Some(EntryOrder::OldestFirst),
2138 ..Default::default()
2139 },
2140 &BranchBounds::default(),
2141 )
2142 .await
2143 .unwrap_or_default();
2144 estimate_context_tokens(&build_session_context(&entries).messages)
2145 }
2146 Err(_) => estimate_context_tokens(&[]),
2150 };
2151 let window = lane
2152 .get_model()
2153 .await
2154 .ok()
2155 .map(|model| model.context_window)
2156 .unwrap_or(0);
2157 let settings = harness.compaction_settings();
2158 let window = window as i64;
2159 let threshold = window - settings.reserve_tokens;
2160 let percent = if window > 0 {
2161 (estimate.tokens as f64) * 100.0 / (window as f64)
2162 } else {
2163 0.0
2164 };
2165 serde_json::json!({
2166 "tokens": estimate.tokens,
2167 "window": window,
2168 "threshold": threshold,
2169 "percent": (percent * 10.0).round() / 10.0,
2170 "reserveTokens": settings.reserve_tokens,
2171 "keepRecentTokens": settings.keep_recent_tokens,
2172 "compactionEnabled": settings.enabled,
2173 "estimated": estimate.last_usage_index.is_none(),
2176 "trailingTokens": estimate.trailing_tokens,
2177 })
2178}
2179
2180fn page_window(
2183 total: usize,
2184 limit: usize,
2185 before: Option<usize>,
2186) -> (std::ops::Range<usize>, bool) {
2187 match before {
2188 Some(index) => {
2189 let start = index.saturating_sub(limit);
2190 (start..index, start > 0)
2191 }
2192 None => {
2193 let skip = total.saturating_sub(limit);
2194 (skip..total, skip > 0)
2195 }
2196 }
2197}
2198
2199async fn rpc_get_messages(
2203 lane: &Arc<dyn AgentLane>,
2204 harness: &AgentHarness,
2205 msg: &serde_json::Value,
2206) -> Result<serde_json::Value, String> {
2207 use rpi_harness::session::types::{BranchBounds, EntryOrder, EntryQuery};
2208
2209 const DEFAULT_LIMIT: usize = 400;
2210 let limit = msg
2211 .get("limit")
2212 .and_then(|value| value.as_u64())
2213 .map(|value| value as usize)
2214 .unwrap_or(DEFAULT_LIMIT)
2215 .max(1);
2216 let before_id = msg
2220 .get("beforeId")
2221 .and_then(|value| value.as_str())
2222 .filter(|value| !value.is_empty())
2223 .map(str::to_owned);
2224 let selected = harness
2225 .transcript_page(&rpi_harness::session::types::TranscriptPageQuery {
2226 limit,
2227 before_id: before_id.clone(),
2228 })
2229 .await
2230 .map_err(|error| error.to_string())?;
2231 let (window, total, has_more, leaf_id) = if let Some(page) = selected {
2232 (
2235 rpc_project_transcript(page.entries)?,
2236 page.total,
2237 page.has_more,
2238 page.leaf_id,
2239 )
2240 } else {
2241 let entries = harness
2243 .find_entries_on_branch(
2244 &EntryQuery {
2245 order: Some(EntryOrder::OldestFirst),
2246 ..Default::default()
2247 },
2248 &BranchBounds::default(),
2249 )
2250 .await
2251 .map_err(|error| error.to_string())?;
2252 let mut items = rpc_project_transcript(entries)?;
2253 let total = items.len();
2254 let before_index = match &before_id {
2255 Some(before) => Some(
2256 items
2257 .iter()
2258 .position(|(id, _)| id == before)
2259 .ok_or_else(|| format!("beforeId not found: {before}"))?,
2260 ),
2261 None => None,
2262 };
2263 let (range, has_more) = page_window(total, limit, before_index);
2264 let window = items.drain(range).collect::<Vec<_>>();
2265 (
2266 window,
2267 total,
2268 has_more,
2269 lane.get_leaf_id().await.ok().flatten(),
2270 )
2271 };
2272 let first_id = window.first().map(|(id, _)| id.clone());
2273 let last_id = window.last().map(|(id, _)| id.clone());
2274 let messages = window
2275 .into_iter()
2276 .map(|(_, value)| value)
2277 .collect::<Vec<_>>();
2278 let pending_resume = lane.has_pending_resume().await;
2279 Ok(serde_json::json!({
2280 "messages": messages,
2281 "total": total,
2282 "truncated": has_more,
2283 "limit": limit,
2284 "firstId": first_id,
2285 "lastId": last_id,
2286 "hasMore": has_more,
2287 "beforeId": before_id,
2288 "leafId": leaf_id,
2289 "pendingResume": pending_resume,
2290 }))
2291}
2292
2293fn rpc_project_transcript(
2294 entries: Vec<rpi_harness::session::types::Entry>,
2295) -> Result<Vec<(String, serde_json::Value)>, String> {
2296 use rpi_harness::session::types::Entry;
2297 let mut items = Vec::new();
2298 for entry in entries {
2299 match entry {
2300 Entry::Message(entry) => {
2301 let mut value = serde_json::to_value(&entry.message).map_err(|e| e.to_string())?;
2302 value["entryId"] = serde_json::json!(entry.base.id);
2304 items.push((entry.base.id, value));
2305 }
2306 Entry::Compaction(compaction) => items.push((
2307 compaction.base.id.clone(),
2308 serde_json::json!({
2309 "role": "compaction",
2310 "entryId": compaction.base.id,
2311 "tokensBefore": compaction.tokens_before,
2312 "retainedMessages": compaction.retained_tail.len(),
2313 "timestamp": compaction.base.timestamp,
2314 }),
2315 )),
2316 _ => {}
2317 }
2318 }
2319 Ok(items)
2320}
2321
2322fn rpc_auth_provider_ids(model_catalog: &[rpi_ai::Model]) -> std::collections::BTreeSet<String> {
2323 let mut providers =
2324 std::collections::BTreeSet::from(["anthropic".to_string(), "openai".to_string()]);
2325 providers.extend(model_catalog.iter().map(|model| model.provider.clone()));
2326 if let Ok(models) = crate::config::load_models_config() {
2327 providers.extend(models.providers.keys().cloned());
2328 }
2329 providers
2330}
2331
2332fn rpc_get_auth_status(
2333 model_catalog: &[rpi_ai::Model],
2334 oauth_snapshot: Option<&Arc<rpi_extensions::RegistrySnapshot>>,
2335) -> serde_json::Value {
2336 let store = crate::config::read_auth().unwrap_or_default();
2337 let models = crate::config::load_models_config().unwrap_or_default();
2338 let manifests = oauth_snapshot.map_or_else(Vec::new, |snapshot| {
2339 snapshot
2340 .providers()
2341 .iter()
2342 .filter_map(|provider| snapshot.provider_manifest(&provider.provider_id).cloned())
2343 .collect()
2344 });
2345 let auth_states = crate::credentials::statuses(&models, &store, &manifests);
2346 let providers = rpc_auth_provider_ids(model_catalog)
2347 .into_iter()
2348 .map(|provider| {
2349 let credential = store.get(&provider);
2350 let state = auth_states.get(&provider).cloned().unwrap_or(
2351 rpi_runtime::credentials::AuthState {
2352 status: rpi_runtime::credentials::AuthStatus::Missing,
2353 source: crate::credentials::CredentialSource::Missing,
2354 },
2355 );
2356 let oauth = matches!(credential, Some(crate::config::Credential::Oauth { .. }));
2357 let oauth_available =
2358 oauth_snapshot.is_some_and(|snapshot| snapshot.oauth_provider(&provider).is_some());
2359 serde_json::json!({
2360 "provider": provider,
2361 "ready": state.is_ready(),
2362 "source": state.source.as_json_str(),
2363 "status": state.status,
2364 "oauth": oauth,
2365 "oauthAvailable": oauth_available,
2366 })
2367 })
2368 .collect::<Vec<_>>();
2369 serde_json::json!({"providers": providers})
2370}
2371
2372fn rpc_oauth(
2373 snapshot: &Option<Arc<rpi_extensions::RegistrySnapshot>>,
2374 msg: &serde_json::Value,
2375) -> Result<serde_json::Value, String> {
2376 let snapshot = snapshot
2377 .as_ref()
2378 .ok_or_else(|| "OAuth is unavailable without an extension provider".to_string())?;
2379 let provider = msg
2380 .get("provider")
2381 .and_then(serde_json::Value::as_str)
2382 .map(str::trim)
2383 .filter(|value| !value.is_empty())
2384 .ok_or_else(|| "oauth requires a provider".to_string())?;
2385 let action = msg
2386 .get("action")
2387 .and_then(serde_json::Value::as_str)
2388 .map(str::trim)
2389 .filter(|value| matches!(*value, "begin" | "exchange" | "refresh" | "revoke"))
2390 .ok_or_else(|| "oauth requires a valid action".to_string())?;
2391 let input = msg
2392 .get("input")
2393 .cloned()
2394 .unwrap_or_else(|| serde_json::json!({}));
2395 let raw = rpi_extensions::request_oauth(snapshot, provider, action, input)
2396 .map_err(|_| "OAuth provider request failed".to_string())?;
2397 let response = rpi_extensions::OAuthActionResponse::parse(raw)
2398 .map_err(|_| "OAuth provider returned an invalid response".to_string())?;
2399 match response.kind {
2400 rpi_extensions::OAuthResponseKind::Credential => {
2401 let credential = response
2402 .credential()
2403 .map_err(|_| "OAuth provider returned an invalid credential".to_string())?
2404 .ok_or_else(|| "OAuth provider returned no credential".to_string())?;
2405 crate::config::upsert_oauth_token(
2406 provider,
2407 credential.access,
2408 credential.refresh,
2409 credential.expires_at,
2410 )
2411 .map_err(|_| "OAuth credentials could not be saved".to_string())?;
2412 Ok(serde_json::json!({
2413 "provider": provider,
2414 "status": "authenticated"
2415 }))
2416 }
2417 rpi_extensions::OAuthResponseKind::Interaction => Ok(serde_json::json!({
2418 "provider": provider,
2419 "status": "interaction",
2420 "display": oauth_display(&response.payload)
2421 })),
2422 rpi_extensions::OAuthResponseKind::Success => Ok(serde_json::json!({
2423 "provider": provider,
2424 "status": "success"
2425 })),
2426 rpi_extensions::OAuthResponseKind::Error => {
2427 Err("OAuth provider rejected the request".to_string())
2428 }
2429 }
2430}
2431
2432fn oauth_display(payload: &serde_json::Value) -> serde_json::Value {
2436 let display = payload.get("display").unwrap_or(payload);
2437 let Some(object) = display.as_object() else {
2438 return serde_json::json!({});
2439 };
2440 let mut safe = serde_json::Map::new();
2441 for key in [
2442 "title",
2443 "url",
2444 "instructions",
2445 "user_code",
2446 "verification_uri",
2447 ] {
2448 if let Some(value) = object.get(key) {
2449 safe.insert(key.to_string(), value.clone());
2450 }
2451 }
2452 serde_json::Value::Object(safe)
2453}
2454
2455async fn rpc_set_api_key(msg: &serde_json::Value) -> Result<serde_json::Value, String> {
2456 let provider = msg
2457 .get("provider")
2458 .and_then(serde_json::Value::as_str)
2459 .map(str::trim)
2460 .filter(|value| !value.is_empty())
2461 .ok_or_else(|| "set_api_key requires a provider".to_string())?;
2462 let api_key = msg
2463 .get("apiKey")
2464 .and_then(serde_json::Value::as_str)
2465 .map(str::trim)
2466 .filter(|value| !value.is_empty())
2467 .ok_or_else(|| "set_api_key requires a non-empty apiKey".to_string())?;
2468 crate::config::upsert_credential(
2469 provider,
2470 crate::config::Credential::ApiKey {
2471 key: Some(api_key.to_string()),
2472 env: None,
2473 },
2474 )
2475 .map_err(|error| error.to_string())?;
2476 Ok(serde_json::json!({
2477 "provider": provider,
2478 "ready": true,
2479 "source": "auth.json",
2480 "requiresRestart": true,
2481 }))
2482}
2483
2484fn rpc_get_available_models(catalog: &[rpi_ai::Model]) -> serde_json::Value {
2485 let models = catalog
2486 .iter()
2487 .map(crate::rpc::protocol::model_descriptor)
2488 .collect::<Vec<_>>();
2489 serde_json::json!({"models": models})
2490}
2491
2492async fn rpc_get_available_thinking_levels(lane: &Arc<dyn AgentLane>) -> serde_json::Value {
2493 let levels = lane
2494 .get_model()
2495 .await
2496 .map(|model| model.supported_thinking_levels())
2497 .unwrap_or_default();
2498 serde_json::json!({"levels": levels})
2499}
2500
2501async fn rpc_configure_external_tools(
2502 lane: &Arc<dyn AgentLane>,
2503 harness: &AgentHarness,
2504 run_active: bool,
2505 backend: &Arc<crate::rpc::external_tools::ExternalTools>,
2506 old_names: &mut Vec<String>,
2507 msg: &serde_json::Value,
2508) -> Result<serde_json::Value, String> {
2509 use rpi_harness::types::HarnessTool;
2510 use rpi_tools::tools::external::{
2511 create_external_tool, validate_external_tool_definitions, ExternalToolDefinition,
2512 };
2513 if run_active
2514 || lane
2515 .snapshot()
2516 .await
2517 .map_err(|e| e.to_string())?
2518 .active
2519 .is_some()
2520 {
2521 return Err("cannot configure external tools while a run is active".into());
2522 }
2523 let definitions: Vec<ExternalToolDefinition> =
2524 serde_json::from_value(msg["tools"].clone()).map_err(|e| e.to_string())?;
2525 let mut registry = harness.get_tools().await.map_err(|e| e.to_string())?;
2526 registry.retain(|tool| !old_names.contains(&tool.tool.schema().name));
2527 let mut reserved = [
2529 "read",
2530 "write",
2531 "edit",
2532 "bash",
2533 "powershell",
2534 "docs",
2535 "grep",
2536 "find",
2537 "ls",
2538 "image",
2539 "permissions",
2540 "permission",
2541 "ask_user",
2542 "plan_mode_start",
2543 "plan_mode_complete",
2544 ]
2545 .iter()
2546 .map(|name| name.to_string())
2547 .collect::<Vec<_>>();
2548 reserved.extend(registry.iter().map(|tool| tool.tool.schema().name.clone()));
2549 validate_external_tool_definitions(&definitions, &reserved).map_err(|e| e.to_string())?;
2550 let names = definitions
2551 .iter()
2552 .map(|def| def.name.clone())
2553 .collect::<Vec<_>>();
2554 let mut active = lane.get_active_tools().await.map_err(|e| e.to_string())?;
2555 active.retain(|name| !old_names.contains(name));
2556 active.extend(names.clone());
2557 registry.extend(
2558 definitions
2559 .into_iter()
2560 .map(|def| HarnessTool::new(create_external_tool(def, backend.clone()))),
2561 );
2562 harness
2563 .set_tools(registry, Some(active.clone()))
2564 .await
2565 .map_err(|e| e.to_string())?;
2566 *old_names = names.clone();
2567 Ok(serde_json::json!({"tools": names, "activeTools": active}))
2568}
2569
2570async fn rpc_get_available_tools(
2571 lane: &Arc<dyn AgentLane>,
2572 harness: &AgentHarness,
2573) -> serde_json::Value {
2574 let active: Vec<String> = if harness.has_exact_resources() {
2575 harness
2576 .execution_context()
2577 .ok()
2578 .and_then(|v| serde_json::from_value(v["tools"].clone()).ok())
2579 .unwrap_or_default()
2580 } else {
2581 lane.get_active_tools().await.unwrap_or_default()
2582 };
2583 let tools = harness
2584 .get_tools()
2585 .await
2586 .map(|tools| {
2587 tools
2588 .iter()
2589 .map(|tool| tool.tool.schema().name.clone())
2590 .collect::<Vec<_>>()
2591 })
2592 .unwrap_or_else(|_| active.clone());
2593 serde_json::json!({"tools": tools, "activeTools": active})
2594}
2595
2596async fn rpc_set_model(
2597 lane: &Arc<dyn AgentLane>,
2598 catalog: &[rpi_ai::Model],
2599 msg: &serde_json::Value,
2600) -> Result<serde_json::Value, String> {
2601 let id = msg
2602 .get("modelId")
2603 .or_else(|| msg.get("model"))
2604 .and_then(|v| v.as_str())
2605 .ok_or("set_model requires `modelId`")?;
2606 let provider = msg.get("provider").and_then(|v| v.as_str());
2607 let model = catalog
2608 .iter()
2609 .find(|m| {
2610 m.id.eq_ignore_ascii_case(id)
2611 && provider.map_or(true, |provider| m.provider.eq_ignore_ascii_case(provider))
2612 })
2613 .cloned()
2614 .ok_or_else(|| format!("unknown model: {id}"))?;
2615 lane.set_model(model).await.map_err(|e| e.to_string())?;
2616 Ok(serde_json::json!({"modelId": id, "provider": provider}))
2617}
2618
2619async fn rpc_set_thinking_level(
2620 lane: &Arc<dyn AgentLane>,
2621 msg: &serde_json::Value,
2622) -> Result<serde_json::Value, String> {
2623 let value = msg
2624 .get("level")
2625 .or_else(|| msg.get("thinkingLevel"))
2626 .ok_or("set_thinking_level requires `level`")?;
2627 let level: rpi_ai::types::ThinkingLevel =
2628 serde_json::from_value(value.clone()).map_err(|_| {
2629 "invalid thinking level (off|minimal|low|medium|high|xhigh|max)".to_string()
2630 })?;
2631 lane.set_thinking_level(level)
2632 .await
2633 .map_err(|e| e.to_string())?;
2634 Ok(serde_json::json!({"thinkingLevel": level}))
2635}
2636
2637async fn rpc_set_active_tools(
2638 lane: &Arc<dyn AgentLane>,
2639 msg: &serde_json::Value,
2640) -> Result<serde_json::Value, String> {
2641 let tools = msg
2642 .get("tools")
2643 .and_then(|v| v.as_array())
2644 .map(|arr| {
2645 arr.iter()
2646 .map(|v| {
2647 v.as_str()
2648 .map(str::to_owned)
2649 .ok_or("tool names must be strings")
2650 })
2651 .collect::<Result<Vec<_>, _>>()
2652 })
2653 .ok_or("set_active_tools requires a `tools` array")??;
2654 lane.set_active_tools(tools.clone())
2655 .await
2656 .map_err(|e| e.to_string())?;
2657 Ok(serde_json::json!({"activeTools": tools}))
2658}
2659
2660fn agent_event_json(event: &AgentEvent) -> serde_json::Value {
2665 crate::rpc::events::project(event)
2666}
2667
2668fn emit_json_event(event: &HarnessEvent) {
2672 let line = match event {
2673 HarnessEvent::CompactionProgress(e) => {
2674 serde_json::json!({ "type": "compaction_progress", "lane": e.lane, "manual": e.manual, "active": e.active })
2675 }
2676 HarnessEvent::RunStart(e) => serde_json::json!({
2677 "type": "run_start",
2678 "lane": e.lane,
2679 "runId": e.run_id,
2680 }),
2681 HarnessEvent::RunEnd(e) => serde_json::json!({
2682 "type": "run_end",
2683 "lane": e.lane,
2684 "runId": e.run_id,
2685 "outcome": run_end_outcome_str(e.outcome),
2686 "leafId": e.leaf_id,
2687 }),
2688 };
2689 println!("{line}");
2690}
2691
2692fn run_end_outcome_str(o: RunEndOutcome) -> &'static str {
2693 match o {
2694 RunEndOutcome::Completed => "completed",
2695 RunEndOutcome::Aborted => "aborted",
2696 RunEndOutcome::Failed => "failed",
2697 }
2698}
2699
2700pub async fn interactive(
2709 harness: &AgentHarness,
2710 event_rx: Option<tokio::sync::broadcast::Receiver<rpi_agent::AgentEvent>>,
2711 args: &Args,
2712 model_catalog: Vec<rpi_ai::Model>,
2713 initial: Option<String>,
2714 extra_messages: &[String],
2715 initial_images: Vec<ImageContent>,
2716 theme: Option<&str>,
2717 no_themes: bool,
2718 reload_context: &crate::session::ReloadContext,
2719) -> i32 {
2720 if let Some(reason) = crate::session::dispatch_session_event_async(
2727 reload_context,
2728 rpi_plugin_sdk::EventTag::BeforeTuiStart,
2729 )
2730 .await
2731 {
2732 eprintln!("[rpi] startup vetoed by extension: {reason}");
2733 return crate::app::EXIT_VETOED;
2734 }
2735
2736 let force_tui = std::env::var("RPI_FORCE_TUI")
2738 .map(|v| v == "1")
2739 .unwrap_or(false);
2740 let code = if force_tui || crate::interactive_tui::is_tui_supported() {
2741 crate::interactive_tui::interactive_tui(
2743 harness,
2744 event_rx,
2745 args,
2746 model_catalog,
2747 initial,
2748 extra_messages,
2749 initial_images,
2750 theme,
2751 no_themes,
2752 reload_context,
2753 )
2754 .await
2755 } else {
2756 for notice in &reload_context.startup_notices {
2759 eprintln!("note: {notice}");
2760 }
2761 interactive_repl(harness, args, initial, extra_messages, initial_images).await
2762 };
2763
2764 if let Some(reason) = crate::session::dispatch_session_event_async(
2770 reload_context,
2771 rpi_plugin_sdk::EventTag::SessionShutdown,
2772 )
2773 .await
2774 {
2775 tracing::warn!(
2776 "[rpi] extension vetoed SessionShutdown (ignored, session closing): {reason}"
2777 );
2778 }
2779
2780 code
2781}
2782
2783pub async fn interactive_repl(
2785 harness: &AgentHarness,
2786 #[allow(unused_variables)] args: &Args,
2787 initial: Option<String>,
2788 extra_messages: &[String],
2789 initial_images: Vec<ImageContent>,
2790) -> i32 {
2791 let lane: Arc<dyn AgentLane> = harness.lane("main");
2793 let stdin = std::io::stdin();
2794 let is_tty = stdin.is_terminal();
2795
2796 if is_tty {
2797 println!(
2798 "rpi interactive (v1 minimal REPL). Type /exit to quit, /abort to cancel a run.\n"
2799 );
2800 }
2801
2802 let mut prompts: Vec<String> = Vec::new();
2804 if let Some(init) = initial {
2805 prompts.push(init);
2806 }
2807 for m in extra_messages {
2808 prompts.push(m.clone());
2809 }
2810 let mut images = initial_images;
2811 for prompt in prompts {
2812 if let Err(code) = run_one(&lane, &prompt, std::mem::take(&mut images)).await {
2813 return code;
2814 }
2815 }
2816
2817 let mut line = String::new();
2819 loop {
2820 if is_tty {
2821 print!("> ");
2822 let _ = std::io::stdout().flush();
2823 }
2824 line.clear();
2825 match stdin.lock().read_line(&mut line) {
2826 Ok(0) => break, Ok(_) => {}
2828 Err(_) => break,
2829 }
2830 let trimmed = line.trim();
2831 if trimmed.is_empty() {
2832 continue;
2833 }
2834 if trimmed == "/exit" || trimmed == "/quit" {
2835 break;
2836 }
2837 if trimmed == "/abort" {
2838 let _ = lane.abort().await;
2839 eprintln!("(aborted)");
2840 continue;
2841 }
2842 if let Err(code) = run_one(&lane, trimmed, Vec::new()).await {
2843 return code;
2844 }
2845 }
2846 0
2847}
2848
2849async fn run_one(
2853 lane: &Arc<dyn AgentLane>,
2854 prompt: &str,
2855 images: Vec<ImageContent>,
2856) -> Result<(), i32> {
2857 match lane.prompt_text(prompt, images).await {
2858 Ok(result) => {
2859 match &result.outcome {
2860 HarnessRunOutcome::Completed { final_message, .. }
2861 | HarnessRunOutcome::Aborted { final_message, .. } => {
2862 let text = assistant_text(final_message);
2863 if !text.is_empty() {
2864 println!("{text}");
2865 }
2866 }
2867 HarnessRunOutcome::Failed {
2868 error,
2869 final_message,
2870 ..
2871 } => {
2872 if let Some(m) = final_message {
2873 if let Some(em) = &m.error_message {
2874 eprintln!("error: {em}");
2875 }
2876 }
2877 eprintln!("run failed: {error:?}");
2878 }
2879 HarnessRunOutcome::Suspended { .. } => {
2880 eprintln!("run suspended (deferred) — resume not supported in v1");
2881 }
2882 }
2883 Ok(())
2884 }
2885 Err(e) => {
2886 eprintln!("prompt rejected: {e}");
2887 Err(1)
2888 }
2889 }
2890}
2891
2892fn session_storage_notice(message: &str) -> String {
2899 format!(
2900 "conversation not saved: {message} — the session log is missing this turn \
2901 and print mode exits without writing it"
2902 )
2903}
2904
2905#[cfg(test)]
2906#[path = "rpc/native_dispatch_tests.rs"]
2907mod native_dispatch_tests;
2908#[cfg(all(test, feature = "experimental-native-store"))]
2909#[path = "modes/native_persistence_tests.rs"]
2910mod native_persistence_tests;
2911#[cfg(test)]
2912#[path = "modes/transcript_page_tests.rs"]
2913mod transcript_page_tests;
2914
2915#[cfg(test)]
2916mod tests {
2917 use super::*;
2918 use rpi_agent::events::AgentEvent;
2919 use rpi_agent::types::AgentToolResult;
2920 use rpi_ai::types::{
2921 AssistantMessage, Content, StopReason, TextContent, TextContentType, Usage,
2922 };
2923 use rpi_harness::session::types::OperationError;
2924
2925 #[test]
2928 fn abort_response_hands_back_the_reclaimed_queues() {
2929 let user = |text: &str| {
2930 rpi_agent::AgentMessage::User(rpi_ai::types::UserMessage::new(
2931 rpi_ai::types::UserContent::Text(text.to_string()),
2932 0,
2933 ))
2934 };
2935 let result = rpi_harness::agent_harness::AbortResult {
2936 run_id: "run-1".to_string(),
2937 steer: vec![user("S")],
2938 follow_up: vec![user("F")],
2939 };
2940
2941 let reply = abort_response(&result);
2942 assert_eq!(reply["aborted"], serde_json::json!(true));
2943 assert_eq!(reply["steer"][0]["text"], serde_json::json!("S"));
2944 assert_eq!(reply["followUp"][0]["text"], serde_json::json!("F"));
2945 assert_eq!(reply["steer"].as_array().unwrap().len(), 1);
2947
2948 assert_eq!(
2950 reclaimed_payload(&result),
2951 serde_json::json!({
2952 "steer": reply["steer"].clone(),
2953 "followUp": reply["followUp"].clone(),
2954 })
2955 );
2956 }
2957
2958 include!("rpc/external_tools_configuration_tests.rs");
2959 include!("rpc/external_tools_lifecycle_tests.rs");
2960 include!("rpc/external_tools_abort_outcome_tests.rs");
2961
2962 #[test]
2963 fn session_branch_query_requires_canonical_arguments() {
2964 let query = branch_entry_query(&serde_json::json!({
2965 "customTypes": ["alpha", "beta", "alpha"], "afterSeq": 42,
2966 }))
2967 .unwrap();
2968 assert_eq!(query.custom_types, ["alpha", "beta"]);
2969 assert_eq!(query.order, EntryOrder::OldestFirst);
2970 assert_eq!(query.after_seq, Some(42));
2971 assert!(query.entry_queries().iter().all(|q| q.limit.is_none()));
2972 for invalid in [
2973 serde_json::json!({}),
2974 serde_json::json!(null),
2975 serde_json::json!([]),
2976 serde_json::json!({"customTypes": []}),
2977 serde_json::json!({"customTypes": "alpha"}),
2978 serde_json::json!({"customTypes": [""]}),
2979 serde_json::json!({"customTypes": [" \t "]}),
2980 serde_json::json!({"customTypes": ["alpha", 1]}),
2981 serde_json::json!({"customTypes": ["alpha"], "limit": 0}),
2982 serde_json::json!({"customTypes": ["alpha"], "limit": 1.5}),
2983 serde_json::json!({"customTypes": ["alpha"], "order": "sideways"}),
2984 serde_json::json!({"customTypes": ["alpha"], "order": 1}),
2985 serde_json::json!({"customTypes": ["alpha"], "afterSeq": -1}),
2986 ] {
2987 assert!(branch_entry_query(&invalid).is_err(), "{invalid}");
2988 }
2989 }
2990
2991 #[test]
2992 fn session_branch_multi_type_limit_is_pushed_to_each_storage_query() {
2993 for order in ["oldestFirst", "newestFirst"] {
2994 let request = branch_entry_query(&serde_json::json!({
2995 "customTypes": ["alpha", "beta", "alpha"],
2996 "limit": 2, "order": order, "afterSeq": 42,
2997 }))
2998 .unwrap();
2999 let queries = request.entry_queries();
3000 assert_eq!(queries.len(), 2);
3001 for query in queries {
3002 assert_eq!(query.limit, Some(2));
3003 assert_eq!(query.order, Some(request.order));
3004 assert_eq!(query.cursor, Some(EntryCursor { after_seq: 42 }));
3005 }
3006 }
3007 }
3008
3009 #[tokio::test]
3010 async fn resource_selection_binding_is_protected_for_scoped_and_unscoped_public_appends() {
3011 use rpi_extensions::RuntimeActionHost;
3012 let harness = AgentHarness::create(Default::default()).await.unwrap();
3013 let (host, cell) = crate::extensions_actions::HarnessActionHost::new_empty(
3014 vec![],
3015 std::env::temp_dir(),
3016 tokio::runtime::Handle::current(),
3017 Default::default(),
3018 rpi_extensions::ToolCallContext::default(),
3019 );
3020 crate::extensions_actions::HarnessActionHost::set_harness(&cell, Arc::new(harness.clone()));
3021 let (_, session_id, leaf_id) = branch_snapshot(
3022 &harness,
3023 &branch_entry_query(&serde_json::json!({"customTypes":["rpi.value"]})).unwrap(),
3024 )
3025 .await
3026 .unwrap();
3027 let guard = serde_json::json!({
3028 "sessionId": session_id, "leafId": leaf_id,
3029 "customType": "rpi.value", "entryId": null
3030 });
3031 for key in ["rpi.resource-selection.v1", "rpi.resource-inventory.v1"] {
3032 for value in [serde_json::Value::Null, serde_json::json!({"forged":true})] {
3033 for scoped in [false, true] {
3034 let mut request = serde_json::json!({
3035 "customType":"rpi.value", "data":{"key":key,"value":value}
3036 });
3037 if scoped {
3038 request["expected"] = guard.clone();
3039 }
3040 assert!(rpc_append_entry(&harness, &request)
3041 .await
3042 .unwrap_err()
3043 .contains("host-owned"));
3044 assert!(host
3045 .append_entry(request)
3046 .await
3047 .unwrap_err()
3048 .contains("host-owned"));
3049 }
3050 }
3051 }
3052 let request = serde_json::json!({
3053 "customType":"rpi.value", "data":{"key":"user.preference","value":1},
3054 "expected":guard
3055 });
3056 assert_eq!(host.append_entry(request).await.unwrap()["scoped"], true);
3057 }
3058
3059 #[tokio::test]
3060 async fn branch_snapshot_guards_and_outbox_acknowledgements_cross_the_runtime_action_boundary()
3061 {
3062 use rpi_extensions::RuntimeActionHost;
3063 let harness = AgentHarness::create(Default::default()).await.unwrap();
3064 let (host, cell) = crate::extensions_actions::HarnessActionHost::new_empty(
3065 vec![],
3066 std::env::temp_dir(),
3067 tokio::runtime::Handle::current(),
3068 Default::default(),
3069 rpi_extensions::ToolCallContext::default(),
3070 );
3071 crate::extensions_actions::HarnessActionHost::set_harness(&cell, Arc::new(harness.clone()));
3072 let before = harness
3074 .session()
3075 .find_entries(&EntryQuery::default())
3076 .await
3077 .unwrap();
3078 let values = host
3079 .get_session_branch(serde_json::json!({"customTypes":["rpi.value"]}))
3080 .await
3081 .unwrap();
3082 for key in ["rpi.resource-selection.v1", "rpi.resource-inventory.v1"] {
3083 let forged = serde_json::json!({
3084 "customType":"rpi.value", "data":{"key":key,"value":{}},
3085 "expected":{"sessionId":values["sessionId"],"leafId":values["leafId"],"customType":"rpi.value","entryId":null}
3086 });
3087 assert!(host
3088 .append_entry(forged.clone())
3089 .await
3090 .unwrap_err()
3091 .contains("host-owned"));
3092 assert!(rpc_append_entry(&harness, &forged)
3093 .await
3094 .unwrap_err()
3095 .contains("host-owned"));
3096 }
3097 assert_eq!(
3098 harness
3099 .session()
3100 .find_entries(&EntryQuery::default())
3101 .await
3102 .unwrap(),
3103 before
3104 );
3105 let snapshot = host
3106 .get_session_branch(
3107 serde_json::json!({"customTypes":["goal-state"],"limit":1,"order":"newestFirst"}),
3108 )
3109 .await
3110 .unwrap();
3111 let expected = serde_json::json!({"sessionId":snapshot["sessionId"],"leafId":snapshot["leafId"],"customType":"goal-state","entryId":null});
3112 let mutation =
3113 serde_json::json!({"customType":"goal-state","data":{"goal":"A"},"expected":expected});
3114 let ack = host.append_entry(mutation.clone()).await.unwrap();
3115 assert_eq!(ack["scoped"], true);
3116 assert!(host.append_entry(mutation.clone()).await.is_err());
3117 assert!(rpc_append_entry(&harness, &mutation).await.is_err());
3118 for _ in 0..2 {
3119 let snapshot = host.get_session_branch(serde_json::json!({"customTypes":["goal-state"],"limit":1,"order":"newestFirst"})).await.unwrap();
3120 let expected = serde_json::json!({"sessionId":snapshot["sessionId"],"leafId":snapshot["leafId"],"customType":"goal-state","entryId":ack["entryId"]});
3121 let notice = host.append_entry(serde_json::json!({"expected":expected,"outboxKey":"r1","message":{
3122 "kind":"custom","role":"custom","content":[{"type":"text","text":"Goal A"}],"data":{"customType":"goal"},"timestamp":1
3123 }})).await.unwrap();
3124 assert_eq!(notice["scoped"], true);
3125 }
3126 let entries = harness
3127 .session()
3128 .find_entries(&EntryQuery {
3129 outbox_key: Some("r1".into()),
3130 ..Default::default()
3131 })
3132 .await
3133 .unwrap();
3134 assert_eq!(entries.len(), 1);
3135 let snapshot = host
3136 .get_session_branch(serde_json::json!({"customTypes":["goal-state"],"limit":1}))
3137 .await
3138 .unwrap();
3139 let expected = serde_json::json!({"sessionId":snapshot["sessionId"],"leafId":snapshot["leafId"],"customType":"goal-state","entryId":ack["entryId"]});
3140 let rpc_ack = rpc_append_entry(
3141 &harness,
3142 &serde_json::json!({"customType":"goal-state","data":{"goal":"B"},"expected":expected}),
3143 )
3144 .await
3145 .unwrap();
3146 assert_eq!(rpc_ack["scoped"], true);
3147 assert!(rpc_ack["entryId"].as_str().is_some_and(|id| !id.is_empty()));
3148 }
3149
3150 #[tokio::test]
3151 async fn session_branch_exact_type_identities_round_trip_rpc_and_actions() {
3152 use rpi_extensions::RuntimeActionHost;
3153 let harness = AgentHarness::create(Default::default()).await.unwrap();
3154 let (host, cell) = crate::extensions_actions::HarnessActionHost::new_empty(
3155 vec![],
3156 std::env::temp_dir(),
3157 tokio::runtime::Handle::current(),
3158 Default::default(),
3159 rpi_extensions::ToolCallContext::default(),
3160 );
3161 crate::extensions_actions::HarnessActionHost::set_harness(&cell, Arc::new(harness.clone()));
3162 let names = ["alpha", " alpha ", "beta", " beta "];
3163 for (index, name) in names.iter().enumerate() {
3164 let args = serde_json::json!({"customType": name, "data": {"n": index}});
3165 if index % 2 == 0 {
3166 host.append_entry(args).await.unwrap();
3167 } else {
3168 rpc_append_entry(&harness, &args).await.unwrap();
3169 }
3170 }
3171 let request = branch_entry_query(&serde_json::json!({
3172 "customTypes": ["alpha", " alpha ", "beta", " beta ", "alpha"],
3173 }))
3174 .unwrap();
3175 assert_eq!(request.custom_types, names);
3176 let mut input = String::new();
3177 let mut expected = Vec::new();
3178 for (index, name) in names.iter().enumerate() {
3179 let args = serde_json::json!({"customTypes": [name]});
3180 let result = host.get_session_branch(args.clone()).await.unwrap();
3181 assert_eq!(result["count"], 1);
3182 assert_eq!(result["entries"][0]["customType"], *name);
3183 assert_eq!(result["entries"][0]["data"]["n"], index);
3184 let mut envelope = args;
3185 envelope["type"] = serde_json::json!("get_session_branch");
3186 envelope["id"] = serde_json::json!(index);
3187 input.push_str(&format!("{envelope}\n"));
3188 expected.push(result);
3189 }
3190 let output = RpcOutput::default();
3191 assert_eq!(
3192 rpc_with_io(
3193 &harness,
3194 None,
3195 vec![],
3196 rpi_extensions::UiDialogMailbox::new(),
3197 None,
3198 crate::args::RpcMessageUpdates::default(),
3199 tokio::io::BufReader::new(input.as_bytes()),
3200 output.clone(),
3201 )
3202 .await,
3203 0
3204 );
3205 let messages = output.messages();
3206 for (index, expected) in expected.into_iter().enumerate() {
3207 let response = messages
3208 .iter()
3209 .find(|message| message["id"] == index)
3210 .unwrap();
3211 assert_eq!(response["data"], expected);
3212 }
3213 }
3214
3215 #[tokio::test]
3216 async fn session_branch_filters_merge_and_respect_fork_bounds() {
3217 use rpi_harness::session::types::{ForkOptions, ForkPosition, SessionRepo};
3218 use rpi_harness::session::{CounterIdGenerator, FakeClock, InMemorySessionRepo, Session};
3219 let repo = InMemorySessionRepo::new(
3220 Arc::new(FakeClock::new()),
3221 Arc::new(CounterIdGenerator::new()),
3222 );
3223 let storage = repo.create(&Default::default()).await.unwrap();
3224 let session = Session::new(Arc::new(storage), None);
3225 let harness = AgentHarness::create(rpi_harness::types::AgentHarnessOptions {
3226 session,
3227 ..Default::default()
3228 })
3229 .await
3230 .unwrap();
3231 let session = harness.session();
3232 let empty = rpc_get_session_branch(
3233 &harness,
3234 &serde_json::json!({"customTypes":["alpha","beta"]}),
3235 )
3236 .await
3237 .unwrap();
3238 assert_eq!(empty["entries"], serde_json::json!([]));
3239 assert_eq!(empty["filtered"], true);
3240 let mut ids = Vec::new();
3241 let mut fork_target = None;
3242 for (index, kind) in ["alpha", "noise", "beta", "alpha", "noise", "beta"]
3243 .into_iter()
3244 .enumerate()
3245 {
3246 if index == 2 {
3247 fork_target = Some(
3248 session
3249 .append_message(rpi_agent::AgentMessage::User(UserMessage::new(
3250 "fork target",
3251 0,
3252 )))
3253 .await
3254 .unwrap(),
3255 );
3256 }
3257 ids.push(session.append_custom_entry(kind, None).await.unwrap());
3258 }
3259 session
3261 .append_message(rpi_agent::AgentMessage::User(UserMessage::new("noise", 0)))
3262 .await
3263 .unwrap();
3264 session
3265 .append_record(
3266 serde_json::from_value(serde_json::json!({
3267 "type":"usage", "id":"usage-noise", "seq":0, "timestamp":0,
3268 "lane":"main", "usage":Usage::zero(), "cause":"adjustment"
3269 }))
3270 .unwrap(),
3271 )
3272 .await
3273 .unwrap();
3274 let run = |args| {
3275 let harness = &harness;
3276 async move { rpc_get_session_branch(harness, &args).await.unwrap() }
3277 };
3278 let all = run(serde_json::json!({"customTypes":["alpha","beta","alpha"]})).await;
3279 let entry_ids = |value: &serde_json::Value| -> Vec<String> {
3280 value["entries"]
3281 .as_array()
3282 .unwrap()
3283 .iter()
3284 .map(|entry| entry["id"].as_str().unwrap().to_owned())
3285 .collect()
3286 };
3287 assert_eq!(
3288 entry_ids(&all),
3289 vec![
3290 ids[0].clone(),
3291 ids[2].clone(),
3292 ids[3].clone(),
3293 ids[5].clone()
3294 ]
3295 );
3296 assert_eq!(all["count"], 4);
3297 assert_eq!(all["filtered"], true);
3298 let cursor = all["entries"][1]["seq"].as_u64().unwrap();
3299 for (order, expected) in [
3300 (
3301 "oldestFirst",
3302 vec![ids[0].clone(), ids[2].clone(), ids[3].clone()],
3303 ),
3304 (
3305 "newestFirst",
3306 vec![ids[5].clone(), ids[3].clone(), ids[2].clone()],
3307 ),
3308 ] {
3309 let result =
3310 run(serde_json::json!({"customTypes":["alpha","beta"], "order":order, "limit":3}))
3311 .await;
3312 assert_eq!(entry_ids(&result), expected);
3313 let result = run(serde_json::json!({"customTypes":["alpha","beta"], "order":order, "afterSeq":cursor, "limit":1})).await;
3314 assert_eq!(
3315 entry_ids(&result),
3316 vec![if order == "oldestFirst" {
3317 ids[3].clone()
3318 } else {
3319 ids[0].clone()
3320 }]
3321 );
3322 }
3323 let single = run(serde_json::json!({"customTypes":["beta"], "afterSeq":all["entries"][3]["seq"], "order":"newestFirst", "limit":1})).await;
3324 assert_eq!(entry_ids(&single), vec![ids[2].clone()]);
3325
3326 session.move_lane("main", Some(&ids[2])).await.unwrap();
3328 let branch_tip = session.append_custom_entry("alpha", None).await.unwrap();
3329 session.create_lane("other", Some(&ids[5])).await.unwrap();
3330 let branch = run(serde_json::json!({"customTypes":["alpha","beta"]})).await;
3331 assert_eq!(
3332 entry_ids(&branch),
3333 vec![ids[0].clone(), ids[2].clone(), branch_tip]
3334 );
3335 let empty =
3336 run(serde_json::json!({"customTypes":["alpha","beta"], "afterSeq":u64::MAX - 1})).await;
3337 assert!(entry_ids(&empty).is_empty());
3338
3339 let fork = repo
3341 .fork(
3342 &session.storage().metadata(),
3343 &Default::default(),
3344 &ForkOptions::Branch {
3345 entry_id: fork_target,
3346 position: Some(ForkPosition::Before),
3347 },
3348 )
3349 .await
3350 .unwrap();
3351 harness
3352 .set_session(Session::new(Arc::new(fork), None))
3353 .await
3354 .unwrap();
3355 let forked = run(
3356 serde_json::json!({"customTypes":["alpha","beta"], "order":"newestFirst", "limit":2}),
3357 )
3358 .await;
3359 assert_eq!(entry_ids(&forked), vec![ids[0].clone()]);
3360 }
3361
3362 #[tokio::test]
3363 async fn session_branch_rpc_and_action_share_strict_query_behavior() {
3364 use rpi_extensions::RuntimeActionHost;
3365 let harness = AgentHarness::create(Default::default()).await.unwrap();
3366 let (host, cell) = crate::extensions_actions::HarnessActionHost::new_empty(
3367 vec![],
3368 std::env::temp_dir(),
3369 tokio::runtime::Handle::current(),
3370 Default::default(),
3371 rpi_extensions::ToolCallContext::default(),
3372 );
3373 crate::extensions_actions::HarnessActionHost::set_harness(&cell, Arc::new(harness.clone()));
3374 host.append_entry(serde_json::json!({"customType":"alpha", "data":{"n":1}}))
3376 .await
3377 .unwrap();
3378 rpc_append_entry(
3379 &harness,
3380 &serde_json::json!({"customType":"beta", "data":{"n":2}}),
3381 )
3382 .await
3383 .unwrap();
3384 let valid =
3385 serde_json::json!({"customTypes":["alpha","beta"], "limit":1, "order":"newestFirst"});
3386 let expected = host.get_session_branch(valid.clone()).await.unwrap();
3387 assert_eq!(expected["entries"][0]["customType"], "beta");
3388 let mut requests = vec![serde_json::json!({})];
3389 for key in [
3390 "customType",
3391 "custom_type",
3392 "custom_types",
3393 "after_seq",
3394 "unknown",
3395 ] {
3396 let value = if key == "after_seq" {
3397 serde_json::json!(0)
3398 } else {
3399 serde_json::json!("alpha")
3400 };
3401 requests.push(serde_json::json!({key:value.clone()}));
3402 let mut mixed = valid.clone();
3403 mixed[key] = value;
3404 requests.push(mixed);
3405 }
3406 let mut input = String::new();
3407 let mut envelope = valid;
3408 envelope["type"] = serde_json::json!("get_session_branch");
3409 envelope["id"] = serde_json::json!("modern");
3410 input.push_str(&format!("{envelope}\n"));
3411 for (index, args) in requests.iter().enumerate() {
3412 let error = host.get_session_branch(args.clone()).await.unwrap_err();
3413 assert!(error.contains("get_session_branch"), "{error}");
3414 let mut envelope = args.clone();
3415 envelope["type"] = serde_json::json!("get_session_branch");
3416 envelope["id"] = serde_json::json!(index);
3417 assert_eq!(
3418 rpc_get_session_branch(&harness, &envelope)
3419 .await
3420 .unwrap_err(),
3421 error
3422 );
3423 input.push_str(&format!("{envelope}\n"));
3424 }
3425 let output = RpcOutput::default();
3426 assert_eq!(
3427 rpc_with_io(
3428 &harness,
3429 None,
3430 vec![],
3431 rpi_extensions::UiDialogMailbox::new(),
3432 None,
3433 crate::args::RpcMessageUpdates::default(),
3434 tokio::io::BufReader::new(input.as_bytes()),
3435 output.clone(),
3436 )
3437 .await,
3438 0
3439 );
3440 let messages = output.messages();
3441 assert_eq!(
3442 messages.iter().find(|m| m["id"] == "modern").unwrap()["data"],
3443 expected
3444 );
3445 for index in 0..requests.len() {
3446 let response = messages.iter().find(|m| m["id"] == index).unwrap();
3447 assert_eq!(response["success"], false, "{response}");
3448 }
3449 }
3450
3451 #[tokio::test]
3452 async fn rpc_commands_list_and_execute_registered_handler() {
3453 extern "C" fn echo(
3454 input: rpi_plugin_sdk::StbStringRef,
3455 out: *mut rpi_plugin_sdk::StbString,
3456 _: *mut std::ffi::c_void,
3457 ) -> i32 {
3458 unsafe {
3460 *out = rpi_plugin_sdk::StbString::from_string(input.as_str().to_owned());
3461 }
3462 0
3463 }
3464 let mut registry = rpi_extensions::ExtensionRegistry::new();
3465 registry.register_command(
3466 "echo".into(),
3467 "Echo arguments".into(),
3468 echo,
3469 std::ptr::null_mut(),
3470 );
3471 let harness = AgentHarness::create(Default::default()).await.unwrap();
3472 let output = RpcOutput::default();
3473 let input = concat!(
3474 "{\"id\":1,\"type\":\"get_commands\"}\n",
3475 "{\"id\":2,\"type\":\"execute_command\",\"command\":\"/echo\",\"args\":\"hello world\"}\n",
3476 "{\"id\":3,\"type\":\"execute_command\",\"command\":\"missing\"}\n",
3477 "{\"id\":4,\"type\":\"execute_command\",\"command\":\"echo\",\"args\":42}\n"
3478 );
3479 assert_eq!(
3480 rpc_with_io(
3481 &harness,
3482 None,
3483 vec![],
3484 rpi_extensions::UiDialogMailbox::new(),
3485 Some(Arc::new(registry.snapshot())),
3486 crate::args::RpcMessageUpdates::default(),
3487 tokio::io::BufReader::new(input.as_bytes()),
3488 output.clone()
3489 )
3490 .await,
3491 0
3492 );
3493 let messages = output.messages();
3494 let response = |id| messages.iter().find(|m| m["id"] == id).unwrap();
3495 assert_eq!(response(1)["data"]["commands"][0]["name"], "echo");
3496 assert_eq!(
3497 response(2)["data"],
3498 serde_json::json!({"command":"/echo", "args":"hello world"})
3499 );
3500 assert_eq!(response(3)["success"], false);
3501 assert_eq!(response(4)["success"], false);
3502 assert_eq!(rpc_get_commands(None), serde_json::json!({"commands":[]}));
3503 }
3504
3505 #[test]
3506 fn rpc_auth_status_shares_header_readiness_and_expired_token_state() {
3507 let _guard = crate::config::test_support::env_lock().lock().unwrap();
3508 struct RestoreConfig(Option<std::ffi::OsString>);
3509 impl Drop for RestoreConfig {
3510 fn drop(&mut self) {
3511 match self.0.take() {
3512 Some(value) => std::env::set_var(crate::config::CONFIG_DIR_ENV, value),
3513 None => std::env::remove_var(crate::config::CONFIG_DIR_ENV),
3514 }
3515 }
3516 }
3517 let _restore = RestoreConfig(std::env::var_os(crate::config::CONFIG_DIR_ENV));
3518 let directory = tempfile::tempdir().unwrap();
3519 std::env::set_var(crate::config::CONFIG_DIR_ENV, directory.path());
3520 std::fs::write(directory.path().join("models.json"), serde_json::json!({"providers": {
3521 "header-gateway": {"api":"openai-completions", "headers":{"Authorization":"Bearer header-secret"}, "models":[{"id":"model"}]},
3522 "expired-account": {"api":"openai-completions", "models":[]}
3523 }}).to_string()).unwrap();
3524 std::fs::write(directory.path().join("auth.json"), serde_json::json!({
3525 "expired-account": {"type":"oauth", "access":"expired-secret", "refresh":"refresh-secret", "expires":1}
3526 }).to_string()).unwrap();
3527 let status = rpc_get_auth_status(&[], None);
3528 let providers = status["providers"].as_array().unwrap();
3529 let header = providers
3530 .iter()
3531 .find(|row| row["provider"] == "header-gateway")
3532 .unwrap();
3533 assert_eq!(header["ready"], true);
3534 assert_eq!(header["source"], "headers");
3535 assert_eq!(header["status"], "ready");
3536 let expired = providers
3537 .iter()
3538 .find(|row| row["provider"] == "expired-account")
3539 .unwrap();
3540 assert_eq!(expired["ready"], false);
3541 assert_eq!(expired["status"], "expired");
3542 assert!(!status.to_string().contains("secret"));
3543 }
3544
3545 #[test]
3546 fn oauth_display_drops_token_like_fields() {
3547 let display = oauth_display(&serde_json::json!({
3548 "display": {
3549 "title": "Sign in",
3550 "url": "https://example.test/authorize",
3551 "instructions": ["Paste the code"],
3552 "access_token": "secret-access",
3553 "refresh_token": "secret-refresh",
3554 }
3555 }));
3556 assert_eq!(display["title"], "Sign in");
3557 assert!(display.get("access_token").is_none());
3558 assert!(display.get("refresh_token").is_none());
3559 }
3560
3561 #[test]
3562 fn rpc_message_text_joins_content_blocks_with_real_newlines() {
3563 let text = rpc_message_text(&serde_json::json!({
3564 "content": [
3565 {"type": "text", "text": "first"},
3566 {"type": "text", "text": "second"}
3567 ]
3568 }));
3569 assert_eq!(text, "first\nsecond");
3570 assert!(!text.contains("\\n"));
3571 }
3572
3573 #[derive(Clone, Default)]
3574 pub(super) struct RpcOutput(Arc<Mutex<Vec<u8>>>);
3575
3576 impl Write for RpcOutput {
3577 fn write(&mut self, bytes: &[u8]) -> std::io::Result<usize> {
3578 self.0.lock().unwrap().extend_from_slice(bytes);
3579 Ok(bytes.len())
3580 }
3581 fn flush(&mut self) -> std::io::Result<()> {
3582 Ok(())
3583 }
3584 }
3585
3586 impl RpcOutput {
3587 pub(super) fn messages(&self) -> Vec<serde_json::Value> {
3588 String::from_utf8(self.0.lock().unwrap().clone())
3589 .unwrap()
3590 .lines()
3591 .map(|line| serde_json::from_str(line).unwrap())
3592 .collect()
3593 }
3594
3595 pub(super) async fn wait_for(
3596 &self,
3597 predicate: impl Fn(&serde_json::Value) -> bool,
3598 ) -> serde_json::Value {
3599 tokio::time::timeout(std::time::Duration::from_secs(3), async {
3600 loop {
3601 if let Some(message) = self.messages().into_iter().find(&predicate) {
3602 return message;
3603 }
3604 tokio::time::sleep(std::time::Duration::from_millis(1)).await;
3605 }
3606 })
3607 .await
3608 .expect("RPC message not received")
3609 }
3610 }
3611
3612 pub(super) async fn rpc_test_harness(
3613 text: &str,
3614 tokens_per_second: f64,
3615 ) -> (AgentHarness, Arc<rpi_ai::providers::faux::FauxProvider>) {
3616 let provider = rpi_ai::providers::faux::FauxProvider::new(
3617 rpi_ai::providers::faux::FauxScript::new()
3618 .with_text(text)
3619 .with_tokens_per_second(tokens_per_second),
3620 );
3621 let options = rpi_harness::types::AgentHarnessOptions {
3622 model: provider.default_model().clone(),
3623 models: vec![provider.clone()],
3624 ..Default::default()
3625 };
3626 (AgentHarness::create(options).await.unwrap(), provider)
3627 }
3628
3629 pub(super) fn start_test_rpc(
3630 harness: AgentHarness,
3631 ui: rpi_extensions::UiDialogMailbox,
3632 ) -> (
3633 tokio::io::DuplexStream,
3634 RpcOutput,
3635 tokio::task::JoinHandle<i32>,
3636 ) {
3637 let (client, input) = tokio::io::duplex(4096);
3638 let output = RpcOutput::default();
3639 let captured = output.clone();
3640 let task = tokio::spawn(async move {
3641 rpc_with_io(
3642 &harness,
3643 None,
3644 vec![],
3645 ui,
3646 None,
3647 crate::args::RpcMessageUpdates::default(),
3648 tokio::io::BufReader::new(input),
3649 captured,
3650 )
3651 .await
3652 });
3653 (client, output, task)
3654 }
3655
3656 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
3657 async fn rpc_paginated_user_summaries_and_context_windows_are_read_only() {
3658 use tokio::io::AsyncWriteExt;
3659 let (harness, _unused_provider) = rpc_test_harness("must not be called", 0.0).await;
3660 let view = harness.view();
3661 for index in 0..23 {
3662 view.append_message(
3663 rpi_ai::UserMessage::new(format!("Question {index}"), index).into(),
3664 )
3665 .await
3666 .unwrap();
3667 }
3668 let pinned = view.get_leaf_id().await.unwrap();
3669 let (mut client, output, task) =
3670 start_test_rpc(harness, rpi_extensions::UiDialogMailbox::new());
3671 let ready = output.wait_for(|frame| frame["type"] == "ready").await;
3672 for command in ["get_message_summaries", "get_message_window"] {
3673 assert!(ready["supportedCommands"]
3674 .as_array()
3675 .unwrap()
3676 .contains(&serde_json::json!(command)));
3677 }
3678 client
3679 .write_all(b"{\"type\":\"get_message_summaries\",\"id\":1,\"limit\":10}\n")
3680 .await
3681 .unwrap();
3682 let first = output.wait_for(|frame| frame["id"] == 1).await;
3683 assert_eq!(first["success"], true);
3684 assert_eq!(first["data"]["messages"][0]["preview"], "Question 13");
3685 assert_eq!(first["data"]["messages"].as_array().unwrap().len(), 10);
3686 for (id, command, cursor, value) in [
3687 (
3688 2,
3689 "get_message_summaries",
3690 "beforeId",
3691 first["data"]["nextBeforeId"].clone(),
3692 ),
3693 (
3694 3,
3695 "get_message_window",
3696 "anchorId",
3697 first["data"]["firstId"].clone(),
3698 ),
3699 ] {
3700 let mut request = serde_json::json!({"id":id,"type":command,"limit":4,"leafId":pinned});
3701 request[cursor] = value;
3702 client
3703 .write_all(format!("{request}\n").as_bytes())
3704 .await
3705 .unwrap();
3706 let page = output.wait_for(|frame| frame["id"] == id).await;
3707 assert_eq!(page["success"], true);
3708 assert_eq!(page["data"]["messages"].as_array().unwrap().len(), 4);
3709 assert_eq!(page["data"]["leafId"], serde_json::json!(pinned));
3710 }
3711 let context = output.wait_for(|frame| frame["id"] == 3).await;
3712 for (id, cursor, key) in [
3713 (4, "beforeId", "nextBeforeId"),
3714 (5, "afterId", "nextAfterId"),
3715 ] {
3716 let mut request =
3717 serde_json::json!({"type":"get_message_window", "id":id,"leafId":pinned,"limit":4});
3718 request[cursor] = context["data"][key].clone();
3719 client
3720 .write_all(format!("{request}\n").as_bytes())
3721 .await
3722 .unwrap();
3723 let response = output.wait_for(|frame| frame["id"] == id).await;
3724 assert_eq!(response["success"], true);
3725 assert_eq!(response["data"]["messages"].as_array().unwrap().len(), 4);
3726 }
3727 client
3728 .write_all(b"{\"type\":\"stop\",\"id\":6}\n")
3729 .await
3730 .unwrap();
3731 assert_eq!(task.await.unwrap(), 0);
3732 assert_eq!(view.get_leaf_id().await.unwrap(), pinned);
3733 assert!(!output
3734 .messages()
3735 .iter()
3736 .any(|frame| frame["type"] == "agent_start"));
3737 }
3738
3739 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
3740 async fn rpc_eof_waits_for_accepted_prompt() {
3741 use tokio::io::AsyncWriteExt;
3742 let (harness, _) = rpc_test_harness("completed", 100.0).await;
3743 let ui = rpi_extensions::UiDialogMailbox::new();
3744 let (mut client, output, task) = start_test_rpc(harness, ui.clone());
3745 client
3746 .write_all(b"{\"type\":\"prompt\",\"id\":1,\"message\":\"test\"}\n")
3747 .await
3748 .unwrap();
3749 client.shutdown().await.unwrap();
3750 assert_eq!(
3751 tokio::time::timeout(std::time::Duration::from_secs(3), task)
3752 .await
3753 .unwrap()
3754 .unwrap(),
3755 0
3756 );
3757 let response = output
3758 .messages()
3759 .into_iter()
3760 .find(|m| m["id"] == 1)
3761 .unwrap();
3762 assert_eq!(response["command"], "prompt");
3763 assert_eq!(response["data"]["accepted"], true);
3764 assert!(output.messages().iter().any(|message| {
3765 message["type"] == "agent_settled" && message["outcome"] == "completed"
3766 }));
3767 assert!(!ui.is_attached());
3768 }
3769
3770 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
3771 async fn rpc_append_entry_publishes_generic_commit_event() {
3772 use tokio::io::AsyncWriteExt;
3773 let (harness, _) = rpc_test_harness("unused", 0.0).await;
3774 let (mut client, output, task) =
3775 start_test_rpc(harness, rpi_extensions::UiDialogMailbox::new());
3776 client
3777 .write_all(
3778 br#"{"type":"append_entry","id":1,"customType":"durable-state","data":{"revision":1}}
3779"#,
3780 )
3781 .await
3782 .unwrap();
3783
3784 let event = output
3785 .wait_for(|message| message["type"] == "session_entry_committed")
3786 .await;
3787 assert_eq!(event["customType"], "durable-state");
3788 assert_eq!(event["entryType"], "custom");
3789 assert_eq!(event["lane"], "main");
3790 assert!(event["entryId"].as_str().is_some());
3791
3792 client
3793 .write_all(
3794 br#"{"type":"stop","id":2}
3795"#,
3796 )
3797 .await
3798 .unwrap();
3799 assert_eq!(task.await.unwrap(), 0);
3800 }
3801
3802 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
3803 async fn rpc_prompt_publishes_committed_message_event() {
3804 use tokio::io::AsyncWriteExt;
3805 let (harness, _) = rpc_test_harness("completed", 0.0).await;
3806 let (mut client, output, task) =
3807 start_test_rpc(harness, rpi_extensions::UiDialogMailbox::new());
3808 client
3809 .write_all(b"{\"type\":\"prompt\",\"id\":1,\"message\":\"test\"}\n")
3810 .await
3811 .unwrap();
3812 let event = output
3813 .wait_for(|message| {
3814 message["type"] == "session_entry_committed" && message["entryType"] == "message"
3815 })
3816 .await;
3817 assert_eq!(event["lane"], "main");
3818 assert!(event["entryId"].as_str().is_some());
3819 output
3820 .wait_for(|message| message["type"] == "agent_settled")
3821 .await;
3822 client
3823 .write_all(b"{\"type\":\"stop\",\"id\":2}\n")
3824 .await
3825 .unwrap();
3826 assert_eq!(task.await.unwrap(), 0);
3827 }
3828
3829 #[tokio::test]
3830 async fn rpc_stop_flushes_all_accepted_entry_invalidations() {
3831 let (harness, _) = rpc_test_harness("unused", 0.0).await;
3832 let output = RpcOutput::default();
3833 let mut input = String::new();
3834 for id in 1..=4 {
3835 input.push_str(&format!(
3836 "{{\"type\":\"append_entry\",\"id\":{id},\"customType\":\"durable-state\"}}\n"
3837 ));
3838 }
3839 input.push_str("{\"type\":\"stop\",\"id\":5}\n");
3840 let result = rpc_with_io(
3841 &harness,
3842 None,
3843 vec![],
3844 rpi_extensions::UiDialogMailbox::new(),
3845 None,
3846 crate::args::RpcMessageUpdates::default(),
3847 tokio::io::BufReader::new(input.as_bytes()),
3848 output.clone(),
3849 )
3850 .await;
3851 assert_eq!(result, 0);
3852 let messages = output.messages();
3853 let commits: std::collections::HashSet<_> = messages
3854 .iter()
3855 .filter(|m| m["type"] == "session_entry_committed")
3856 .map(|m| m["entryId"].as_str().unwrap())
3857 .collect();
3858 assert_eq!(
3859 commits.len(),
3860 4,
3861 "stop must flush every committed entry: {messages:?}"
3862 );
3863 }
3864
3865 #[tokio::test]
3866 async fn rpc_eof_flushes_queued_entry_invalidations() {
3867 let (harness, _) = rpc_test_harness("unused", 0.0).await;
3868 let output = RpcOutput::default();
3869 let input = b"{\"type\":\"append_entry\",\"id\":1,\"customType\":\"durable-state\",\"data\":{\"revision\":1}}\n";
3870 let result = rpc_with_io(
3871 &harness,
3872 None,
3873 vec![],
3874 rpi_extensions::UiDialogMailbox::new(),
3875 None,
3876 crate::args::RpcMessageUpdates::default(),
3877 tokio::io::BufReader::new(&input[..]),
3878 output.clone(),
3879 )
3880 .await;
3881 assert_eq!(result, 0);
3882 let messages = output.messages();
3883 assert!(
3884 messages
3885 .iter()
3886 .any(|m| m["type"] == "response" && m["id"] == 1 && m["success"] == true),
3887 "append succeeded: {messages:?}"
3888 );
3889 assert!(
3890 messages
3891 .iter()
3892 .any(|m| m["type"] == "session_entry_committed"),
3893 "EOF must flush the queued commit event: {messages:?}"
3894 );
3895 }
3896
3897 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
3898 async fn rpc_compaction_settlement_identifies_its_durable_operation() {
3899 let (harness, _) = rpc_test_harness("unused", 0.0).await;
3900 let output = RpcOutput::default();
3901 let code = rpc_with_io(
3902 &harness,
3903 None,
3904 vec![],
3905 rpi_extensions::UiDialogMailbox::new(),
3906 None,
3907 crate::args::RpcMessageUpdates::default(),
3908 tokio::io::BufReader::new(&b"{\"type\":\"compact\",\"id\":1}\n"[..]),
3909 output.clone(),
3910 )
3911 .await;
3912 assert_eq!(code, 0);
3913 let messages = output.messages();
3914 let settled = messages
3915 .iter()
3916 .find(|message| message["type"] == "agent_settled")
3917 .unwrap();
3918 assert_eq!(settled["outcome"], "declined");
3919 let run_id = settled["runId"]
3920 .as_str()
3921 .expect("compaction settlement needs its durable run id");
3922 let records = harness
3923 .session()
3924 .find_records(&Default::default())
3925 .await
3926 .unwrap();
3927 assert!(records.iter().any(|record| matches!(record,
3928 rpi_harness::session::types::LaneRecord::OperationFinished(finished)
3929 if finished.run_id == run_id)));
3930 }
3931
3932 #[tokio::test]
3933 async fn rpc_eof_drains_queued_final_events() {
3934 let harness = AgentHarness::create(Default::default()).await.unwrap();
3935 let (tx, rx) = tokio::sync::broadcast::channel(256);
3936 for _ in 0..100 {
3937 tx.send(AgentEvent::AgentStart).unwrap();
3938 }
3939 tx.send(AgentEvent::AgentEnd { messages: vec![] }).unwrap();
3940 let output = RpcOutput::default();
3941 let code = rpc_with_io(
3942 &harness,
3943 Some(rx),
3944 vec![],
3945 rpi_extensions::UiDialogMailbox::new(),
3946 None,
3947 crate::args::RpcMessageUpdates::default(),
3948 tokio::io::BufReader::new(&b""[..]),
3949 output.clone(),
3950 )
3951 .await;
3952 assert_eq!(code, 0);
3953 let messages = output.messages();
3954 assert_eq!(
3955 messages.len(),
3956 102,
3957 "ready plus every queued event must be flushed"
3958 );
3959 assert_eq!(messages.last().unwrap()["type"], "agent_end");
3960 }
3961
3962 fn assert_rpc_update_mode(
3963 messages: &[serde_json::Value],
3964 mode: crate::args::RpcMessageUpdates,
3965 ) {
3966 let updates: Vec<_> = messages
3967 .iter()
3968 .filter(|m| {
3969 matches!(
3970 m["eventType"].as_str(),
3971 Some("text_delta" | "thinking_delta" | "toolcall_delta")
3972 )
3973 })
3974 .collect();
3975 assert!(
3976 !updates.is_empty(),
3977 "expected streaming deltas: {messages:?}"
3978 );
3979 for update in updates {
3980 let snapshot = mode == crate::args::RpcMessageUpdates::Snapshot;
3981 assert_eq!(update.get("message").is_some(), snapshot, "{update}");
3982 assert_eq!(
3983 update["assistantMessageEvent"].get("partial").is_some(),
3984 snapshot,
3985 "{update}"
3986 );
3987 assert_eq!(update["eventType"], update["assistantMessageEvent"]["type"]);
3988 assert_eq!(
3989 update["contentIndex"],
3990 update["assistantMessageEvent"]["contentIndex"]
3991 );
3992 assert_eq!(update["delta"], update["assistantMessageEvent"]["delta"]);
3993 }
3994 }
3995
3996 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
3997 async fn rpc_selected_updates_survive_prompt_eof_and_stop() {
3998 use tokio::io::AsyncWriteExt;
3999 for mode in [
4000 crate::args::RpcMessageUpdates::CompactDeltas,
4001 crate::args::RpcMessageUpdates::Snapshot,
4002 ] {
4003 for stop in [false, true] {
4004 let (harness, _) =
4005 rpc_test_harness("one two three four five six seven eight", 100.0).await;
4006 let (emitter, events) = rpi_agent::events::BroadcastEmitter::new(256);
4007 harness
4008 .set_agent_emitter(Some(Arc::new(emitter)))
4009 .await
4010 .unwrap();
4011 let (mut client, input) = tokio::io::duplex(4096);
4012 let output = RpcOutput::default();
4013 let captured = output.clone();
4014 let task = tokio::spawn(async move {
4015 rpc_with_io(
4016 &harness,
4017 Some(events),
4018 vec![],
4019 rpi_extensions::UiDialogMailbox::new(),
4020 None,
4021 mode,
4022 tokio::io::BufReader::new(input),
4023 captured,
4024 )
4025 .await
4026 });
4027 let ready = output.wait_for(|m| m["type"] == "ready").await;
4028 assert_eq!(ready["semantics"]["messageUpdates"], mode.as_str());
4029 client
4030 .write_all(b"{\"type\":\"prompt\",\"id\":1,\"message\":\"test\"}\n")
4031 .await
4032 .unwrap();
4033 if stop {
4034 output.wait_for(|m| m["eventType"] == "text_delta").await;
4035 client
4036 .write_all(b"{\"type\":\"stop\",\"id\":2}\n")
4037 .await
4038 .unwrap();
4039 } else {
4040 client.shutdown().await.unwrap();
4041 }
4042 assert_eq!(
4043 tokio::time::timeout(std::time::Duration::from_secs(3), task)
4044 .await
4045 .unwrap()
4046 .unwrap(),
4047 0
4048 );
4049 let messages = output.messages();
4050 assert_rpc_update_mode(&messages, mode);
4051 assert!(messages.iter().any(|m| m["type"] == "message_end"));
4052 assert!(messages.iter().any(|m| m["type"] == "agent_end"));
4053 assert!(messages.iter().any(|m| m["type"] == "agent_settled"));
4054 if !stop {
4055 assert!(messages
4056 .iter()
4057 .any(|m| m["type"] == "agent_settled" && m["outcome"] == "completed"));
4058 }
4059 }
4060 }
4061 }
4062
4063 #[tokio::test]
4064 async fn rpc_normal_and_explicit_drain_use_the_same_immutable_writer_mode() {
4065 for mode in [
4066 crate::args::RpcMessageUpdates::CompactDeltas,
4067 crate::args::RpcMessageUpdates::Snapshot,
4068 ] {
4069 let output = RpcOutput::default();
4070 let (writer, task, _) = start_rpc_writer(output.clone(), mode);
4071 let cloned_writer = writer.clone();
4072 let mut message = rpi_ai::AssistantMessage::empty(rpi_ai::Api::Faux, "faux", "m", 42);
4073 message.content = vec![rpi_ai::Content::text("snapshot")];
4074 let message = Arc::new(message);
4075 let (tx, mut rx) = tokio::sync::broadcast::channel(16);
4076 for kind in 0..3 {
4077 let assistant_message_event = match kind {
4078 0 => rpi_ai::AssistantMessageEvent::TextDelta {
4079 content_index: 0,
4080 delta: "text".into(),
4081 partial: message.clone(),
4082 },
4083 1 => rpi_ai::AssistantMessageEvent::ThinkingDelta {
4084 content_index: 1,
4085 delta: "thought".into(),
4086 partial: message.clone(),
4087 },
4088 _ => rpi_ai::AssistantMessageEvent::ToolCallDelta {
4089 content_index: 2,
4090 delta: "{}".into(),
4091 partial: message.clone(),
4092 },
4093 };
4094 let event = AgentEvent::MessageUpdate {
4095 message_id: None,
4096 message: message.clone(),
4097 assistant_message_event,
4098 };
4099 write_agent_event(&writer, &event).await;
4100 tx.send(event).unwrap();
4101 }
4102 drain_rpc_events(&cloned_writer, &mut rx).await;
4103 drop(cloned_writer);
4104 drop(writer);
4105 task.await.unwrap().unwrap();
4106 let messages = output.messages();
4107 assert_eq!(messages.len(), 6);
4108 assert_eq!(messages[..3], messages[3..]);
4109 assert_rpc_update_mode(&messages, mode);
4110 }
4111 }
4112
4113 #[tokio::test]
4114 async fn rpc_queued_deltas_are_flushed_on_eof_and_stop_in_selected_mode() {
4115 for mode in [
4116 crate::args::RpcMessageUpdates::CompactDeltas,
4117 crate::args::RpcMessageUpdates::Snapshot,
4118 ] {
4119 for input in ["", "{\"type\":\"stop\",\"id\":1}\n"] {
4120 let harness = AgentHarness::create(Default::default()).await.unwrap();
4121 let (tx, rx) = tokio::sync::broadcast::channel(256);
4122 let mut message =
4123 rpi_ai::AssistantMessage::empty(rpi_ai::Api::Faux, "faux", "m", 42);
4124 message.content = vec![rpi_ai::Content::text("snapshot")];
4125 let message = Arc::new(message);
4126 for _ in 0..100 {
4127 tx.send(AgentEvent::MessageUpdate {
4128 message_id: None,
4129 message: message.clone(),
4130 assistant_message_event: rpi_ai::AssistantMessageEvent::TextDelta {
4131 content_index: 0,
4132 delta: "x".into(),
4133 partial: message.clone(),
4134 },
4135 })
4136 .unwrap();
4137 }
4138 tx.send(AgentEvent::AgentEnd { messages: vec![] }).unwrap();
4139 let output = RpcOutput::default();
4140 assert_eq!(
4141 rpc_with_io(
4142 &harness,
4143 Some(rx),
4144 vec![],
4145 rpi_extensions::UiDialogMailbox::new(),
4146 None,
4147 mode,
4148 tokio::io::BufReader::new(input.as_bytes()),
4149 output.clone()
4150 )
4151 .await,
4152 0
4153 );
4154 let messages = output.messages();
4155 assert_eq!(
4156 messages
4157 .iter()
4158 .filter(|m| m["eventType"] == "text_delta")
4159 .count(),
4160 100
4161 );
4162 assert_eq!(
4163 messages.iter().find(|m| m["type"] == "ready").unwrap()["semantics"]
4164 ["messageUpdates"],
4165 mode.as_str()
4166 );
4167 assert_rpc_update_mode(&messages, mode);
4168 assert!(messages.iter().any(|m| m["type"] == "agent_end"));
4169 }
4170 }
4171 }
4172
4173 struct BackpressuredOutput {
4174 entered: Arc<std::sync::atomic::AtomicBool>,
4175 release: Arc<std::sync::atomic::AtomicBool>,
4176 output: RpcOutput,
4177 }
4178
4179 impl Write for BackpressuredOutput {
4180 fn write(&mut self, bytes: &[u8]) -> std::io::Result<usize> {
4181 self.entered
4182 .store(true, std::sync::atomic::Ordering::SeqCst);
4183 while !self.release.load(std::sync::atomic::Ordering::SeqCst) {
4184 std::thread::sleep(std::time::Duration::from_millis(1));
4185 }
4186 self.output.write(bytes)
4187 }
4188 fn flush(&mut self) -> std::io::Result<()> {
4189 Ok(())
4190 }
4191 }
4192
4193 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
4194 async fn rpc_output_backpressure_is_bounded_and_keeps_workers_responsive() {
4195 use std::sync::atomic::{AtomicBool, Ordering};
4196 let release = Arc::new(AtomicBool::new(false));
4197 struct ReleaseOnDrop(Arc<AtomicBool>);
4198 impl Drop for ReleaseOnDrop {
4199 fn drop(&mut self) {
4200 self.0.store(true, Ordering::SeqCst);
4201 }
4202 }
4203 let _release = ReleaseOnDrop(release.clone());
4204 let entered = Arc::new(AtomicBool::new(false));
4205 let output = RpcOutput::default();
4206 let (writer, task, _) = start_rpc_writer(
4207 BackpressuredOutput {
4208 entered: entered.clone(),
4209 release: release.clone(),
4210 output: output.clone(),
4211 },
4212 crate::args::RpcMessageUpdates::default(),
4213 );
4214 write_line(&writer, &serde_json::json!({"id": 0})).await;
4215 tokio::time::timeout(std::time::Duration::from_secs(2), async {
4216 while !entered.load(Ordering::SeqCst) {
4217 tokio::task::yield_now().await;
4218 }
4219 })
4220 .await
4221 .unwrap();
4222 let producer_writer = writer.clone();
4223 let producer = tokio::spawn(async move {
4224 for id in 1..=100 {
4225 write_line(&producer_writer, &serde_json::json!({"id": id})).await;
4226 }
4227 });
4228 tokio::time::timeout(
4229 std::time::Duration::from_millis(150),
4230 tokio::time::sleep(std::time::Duration::from_millis(20)),
4231 )
4232 .await
4233 .unwrap();
4234 assert!(
4235 !producer.is_finished(),
4236 "the bounded queue must apply backpressure"
4237 );
4238 release.store(true, Ordering::SeqCst);
4239 producer.await.unwrap();
4240 drop(writer);
4241 task.await.unwrap().unwrap();
4242 let messages = output.messages();
4243 assert_eq!(messages.len(), 101);
4244 for (index, message) in messages.iter().enumerate() {
4245 assert_eq!(message["id"], index);
4246 }
4247 }
4248
4249 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
4250 async fn rpc_completion_response_follows_final_event_under_backpressure() {
4251 use std::sync::atomic::{AtomicBool, Ordering};
4252 let (harness, _) = rpc_test_harness("completed", 0.0).await;
4253 let (emitter, events) = rpi_agent::events::BroadcastEmitter::new(256);
4254 harness
4255 .set_agent_emitter(Some(Arc::new(emitter)))
4256 .await
4257 .unwrap();
4258 let output = RpcOutput::default();
4259 let release = Arc::new(AtomicBool::new(false));
4260 let release_writer = release.clone();
4261 let release_thread = std::thread::spawn(move || {
4262 std::thread::sleep(std::time::Duration::from_millis(50));
4263 release_writer.store(true, Ordering::SeqCst);
4264 });
4265 let code = rpc_with_io(
4266 &harness,
4267 Some(events),
4268 vec![],
4269 rpi_extensions::UiDialogMailbox::new(),
4270 None,
4271 crate::args::RpcMessageUpdates::default(),
4272 tokio::io::BufReader::new(
4273 &b"{\"type\":\"prompt\",\"id\":1,\"message\":\"test\"}\n"[..],
4274 ),
4275 BackpressuredOutput {
4276 entered: Arc::new(AtomicBool::new(false)),
4277 release,
4278 output: output.clone(),
4279 },
4280 )
4281 .await;
4282 release_thread.join().unwrap();
4283 assert_eq!(code, 0);
4284 let messages = output.messages();
4285 let end = messages
4286 .iter()
4287 .position(|message| message["type"] == "agent_end")
4288 .unwrap();
4289 let settled = messages
4290 .iter()
4291 .position(|message| message["type"] == "agent_settled")
4292 .unwrap();
4293 assert!(end < settled);
4294 let run_id = messages[settled]["runId"].as_str().unwrap();
4295 assert!(!run_id.is_empty(), "settlement identifies its durable run");
4296 let records = harness
4297 .session()
4298 .find_records(&Default::default())
4299 .await
4300 .unwrap();
4301 assert!(records.iter().any(|record| matches!(record,
4302 rpi_harness::session::types::LaneRecord::OperationFinished(finished)
4303 if finished.run_id == run_id)));
4304 assert!(messages.iter().any(|message| {
4305 message["type"] == "response" && message["data"]["accepted"] == true
4306 }));
4307 assert!(messages.iter().any(|message| {
4308 message["type"] == "agent_settled" && message["outcome"] == "completed"
4309 }));
4310 }
4311
4312 #[tokio::test]
4313 async fn rpc_output_failure_closes_without_waiting_for_more_input() {
4314 struct BrokenOutput;
4315 impl Write for BrokenOutput {
4316 fn write(&mut self, _: &[u8]) -> std::io::Result<usize> {
4317 Err(std::io::Error::new(
4318 std::io::ErrorKind::BrokenPipe,
4319 "closed client output",
4320 ))
4321 }
4322 fn flush(&mut self) -> std::io::Result<()> {
4323 Ok(())
4324 }
4325 }
4326 let harness = AgentHarness::create(Default::default()).await.unwrap();
4327 let (_client, input) = tokio::io::duplex(4096);
4328 let ui = rpi_extensions::UiDialogMailbox::new();
4329 let code = tokio::time::timeout(
4330 std::time::Duration::from_secs(2),
4331 rpc_with_io(
4332 &harness,
4333 None,
4334 vec![],
4335 ui.clone(),
4336 None,
4337 crate::args::RpcMessageUpdates::default(),
4338 tokio::io::BufReader::new(input),
4339 BrokenOutput,
4340 ),
4341 )
4342 .await
4343 .unwrap();
4344 assert_eq!(code, 1);
4345 assert!(!ui.is_attached());
4346 }
4347
4348 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
4349 async fn rpc_stop_aborts_and_joins_active_prompt() {
4350 use tokio::io::AsyncWriteExt;
4351 let (harness, provider) = rpc_test_harness("still streaming the response", 20.0).await;
4352 let lane = harness.lane("main");
4353 let (mut client, output, task) =
4354 start_test_rpc(harness, rpi_extensions::UiDialogMailbox::new());
4355 client
4356 .write_all(b"{\"type\":\"prompt\",\"id\":1,\"message\":\"test\"}\n")
4357 .await
4358 .unwrap();
4359 tokio::time::timeout(std::time::Duration::from_secs(3), async {
4360 while provider
4361 .state()
4362 .call_count
4363 .load(std::sync::atomic::Ordering::SeqCst)
4364 == 0
4365 {
4366 tokio::task::yield_now().await;
4367 }
4368 })
4369 .await
4370 .unwrap();
4371 client
4372 .write_all(b"{\"type\":\"stop\",\"id\":2}\n")
4373 .await
4374 .unwrap();
4375 assert_eq!(
4376 tokio::time::timeout(std::time::Duration::from_secs(3), task)
4377 .await
4378 .unwrap()
4379 .unwrap(),
4380 0
4381 );
4382 let messages = output.messages();
4383 assert!(messages.iter().any(|message| {
4384 message["type"] == "agent_settled"
4385 && (message["outcome"] == "aborted" || message["outcome"] == "failed")
4386 }));
4387 assert_eq!(
4388 messages.iter().find(|m| m["id"] == 2).unwrap()["data"]["stopped"],
4389 true
4390 );
4391 assert!(
4392 lane.abort().await.is_err(),
4393 "RPC returned with an active lane"
4394 );
4395 }
4396
4397 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
4398 async fn rpc_immediate_stop_cancels_prompt_before_or_after_lane_start() {
4399 use tokio::io::AsyncWriteExt;
4400 let (harness, _) = rpc_test_harness("not completed", 20.0).await;
4401 let lane = harness.lane("main");
4402 let (mut client, output, task) =
4403 start_test_rpc(harness, rpi_extensions::UiDialogMailbox::new());
4404 client.write_all(b"{\"type\":\"prompt\",\"id\":1,\"message\":\"test\"}\n{\"type\":\"stop\",\"id\":2}\n").await.unwrap();
4405 assert_eq!(
4406 tokio::time::timeout(std::time::Duration::from_secs(3), task)
4407 .await
4408 .unwrap()
4409 .unwrap(),
4410 0
4411 );
4412 let messages = output.messages();
4413 let response = messages
4414 .iter()
4415 .find(|message| message["type"] == "response" && message["id"] == 1)
4416 .expect("the prompt has one correlated admission outcome");
4417 assert_eq!(
4418 messages
4419 .iter()
4420 .filter(|message| { message["type"] == "response" && message["id"] == 1 })
4421 .count(),
4422 1
4423 );
4424 if response["success"] == true {
4425 assert_eq!(response["data"]["accepted"], true);
4426 assert!(messages.iter().any(|message| {
4427 message["type"] == "agent_settled"
4428 && (message["outcome"] == "aborted" || message["outcome"] == "failed")
4429 }));
4430 } else {
4431 assert_eq!(
4434 response["error"].as_str(),
4435 Some("Operation aborted during prompt validation"),
4436 "{response}"
4437 );
4438 assert!(!messages
4439 .iter()
4440 .any(|message| message["type"] == "agent_start"));
4441 }
4442 assert!(messages.iter().any(|message| message["type"] == "response"
4443 && message["id"] == 2
4444 && message["success"] == true));
4445 assert!(lane.abort().await.is_err());
4446 }
4447
4448 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
4449 async fn rpc_dialog_answers_and_ping_are_read_while_prompt_runs() {
4450 use tokio::io::AsyncWriteExt;
4451 let (harness, _) =
4452 rpc_test_harness("the streamed response finishes after the dialog", 30.0).await;
4453 let ui = rpi_extensions::UiDialogMailbox::new();
4454 let (mut client, output, task) = start_test_rpc(harness, ui.clone());
4455 output.wait_for(|m| m["type"] == "ready").await;
4456 client
4457 .write_all(b"{\"type\":\"prompt\",\"id\":1,\"message\":\"test\"}\n")
4458 .await
4459 .unwrap();
4460 ui.handle(serde_json::json!({"requestId":"approval", "toolCallId":"call", "ui":{}}))
4461 .unwrap();
4462 output.wait_for(|m| m["type"] == "ui_dialog_request").await;
4463 client.write_all(b"{\"type\":\"ui_dialog_response\",\"requestId\":\"approval\",\"answer\":true}\n{\"type\":\"ping\",\"id\":2}\n").await.unwrap();
4464 output.wait_for(|m| m["id"] == 2).await;
4465 assert_eq!(ui.poll("approval").unwrap()["answer"], true);
4466 output
4467 .wait_for(|m| m["type"] == "ui_dialog_closed" && m["requestId"] == "approval")
4468 .await;
4469 client.shutdown().await.unwrap();
4470 assert_eq!(
4471 tokio::time::timeout(std::time::Duration::from_secs(3), task)
4472 .await
4473 .unwrap()
4474 .unwrap(),
4475 0
4476 );
4477 assert!(output.messages().iter().any(|message| {
4478 message["type"] == "agent_settled" && message["outcome"] == "completed"
4479 }));
4480 }
4481
4482 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
4483 async fn rpc_notifies_client_when_an_approval_is_cancelled() {
4484 use tokio::io::AsyncWriteExt;
4485 let harness = AgentHarness::create(Default::default()).await.unwrap();
4486 let ui = rpi_extensions::UiDialogMailbox::new();
4487 let (mut client, output, task) = start_test_rpc(harness, ui.clone());
4488 output.wait_for(|m| m["type"] == "ready").await;
4489 ui.handle(serde_json::json!({"requestId":"cancelled", "toolCallId":"c", "ui":{}}))
4490 .unwrap();
4491 output.wait_for(|m| m["type"] == "ui_dialog_request").await;
4492 ui.cancel_tool_call("c");
4493 output
4494 .wait_for(|m| m["type"] == "ui_dialog_closed" && m["requestId"] == "cancelled")
4495 .await;
4496 assert_eq!(ui.poll("cancelled").unwrap()["status"], "cancelled");
4498 client.shutdown().await.unwrap();
4499 assert_eq!(task.await.unwrap(), 0);
4500 }
4501
4502 #[tokio::test]
4503 async fn rpc_accepts_native_dialog_responses_without_consuming_plugin_answers() {
4504 use crate::remote::protocol::RemoteCommand;
4505 use tokio::io::AsyncWriteExt;
4506 let harness = AgentHarness::create(Default::default()).await.unwrap();
4507 let ui = rpi_extensions::UiDialogMailbox::new();
4508 let (mut client, output, task) = start_test_rpc(harness, ui.clone());
4509 output.wait_for(|m| m["type"] == "ready").await;
4510 for (request_id, answer, response_id) in [
4511 (
4512 "answered",
4513 Some(serde_json::json!({"values":["No"],"kind":"selector"})),
4514 5,
4515 ),
4516 ("cancelled", None, 6),
4517 ] {
4518 ui.handle(serde_json::json!({"requestId":request_id,"toolCallId":request_id,"ui":{}}))
4519 .unwrap();
4520 output
4521 .wait_for(|m| m["type"] == "ui_dialog_request" && m["requestId"] == request_id)
4522 .await;
4523 let command = RemoteCommand::UiDialogResponse {
4524 request_id: request_id.into(),
4525 cancelled: answer.is_none(),
4526 answer,
4527 }
4528 .with_id(response_id);
4529 client
4530 .write_all(format!("{command}\n").as_bytes())
4531 .await
4532 .unwrap();
4533 let ack = output.wait_for(|m| m["id"] == response_id).await;
4534 assert_eq!(ack["data"]["accepted"], true);
4535 let result = ui.poll(request_id).unwrap();
4536 if request_id == "answered" {
4537 assert_eq!(result["answer"]["values"][0], "No");
4538 } else {
4539 assert_eq!(result["status"], "cancelled");
4540 }
4541 }
4542 client.shutdown().await.unwrap();
4543 assert_eq!(task.await.unwrap(), 0);
4544 }
4545
4546 fn assistant(text: &str, stop: StopReason) -> AssistantMessage {
4547 AssistantMessage {
4548 role: rpi_ai::types::AssistantRole,
4549 content: vec![Content::Text(TextContent {
4550 kind: TextContentType,
4551 text: text.into(),
4552 text_signature: None,
4553 })],
4554 api: rpi_ai::Api::AnthropicMessages,
4555 provider: "anthropic".into(),
4556 model: "claude-sonnet-5".into(),
4557 response_model: None,
4558 response_id: None,
4559 usage: Usage::zero(),
4560 stop_reason: stop,
4561 deferred: None,
4562 error_message: None,
4563 raw_stop_reason: None,
4564 end_turn: None,
4565 timestamp: 0,
4566 }
4567 }
4568
4569 #[test]
4570 fn assistant_text_concatenates_text_blocks() {
4571 let m = assistant("hello", StopReason::Stop);
4572 assert_eq!(assistant_text(&m), "hello");
4573 }
4574
4575 #[test]
4576 fn outcome_exit_code_maps_failed_aborted_to_1() {
4577 let failed = HarnessRunOutcome::Failed {
4578 leaf_id: "l".into(),
4579 error: OperationError {
4580 code: "boom".into(),
4581 message: "boom".into(),
4582 },
4583 final_entry_id: None,
4584 final_message: None,
4585 };
4586 assert_eq!(outcome_exit_code(&failed), 1);
4587 let completed = HarnessRunOutcome::Completed {
4588 leaf_id: "l".into(),
4589 final_entry_id: "e".into(),
4590 final_message: assistant("ok", StopReason::Stop),
4591 };
4592 assert_eq!(outcome_exit_code(&completed), 0);
4593 }
4594
4595 #[test]
4596 fn run_end_outcome_str_roundtrip() {
4597 assert_eq!(run_end_outcome_str(RunEndOutcome::Completed), "completed");
4598 assert_eq!(run_end_outcome_str(RunEndOutcome::Aborted), "aborted");
4599 assert_eq!(run_end_outcome_str(RunEndOutcome::Failed), "failed");
4600 }
4601
4602 #[test]
4603 fn agent_event_projection_keeps_terminal_and_tool_payloads() {
4604 let end = agent_event_json(&AgentEvent::AgentEnd { messages: vec![] });
4605 assert_eq!(end["type"], "agent_end");
4606 assert_eq!(end["messages"], serde_json::json!([]));
4607
4608 let tool = agent_event_json(&AgentEvent::ToolExecutionEnd {
4609 tool_call_id: "call-1".into(),
4610 tool_name: "read".into(),
4611 result: AgentToolResult::text("hello"),
4612 is_error: false,
4613 });
4614 assert_eq!(tool["type"], "tool_execution_end");
4615 assert_eq!(tool["result"]["content"][0]["text"], "hello");
4616 assert_eq!(tool["result"]["terminate"], false);
4617
4618 let retry = agent_event_json(&AgentEvent::RetryScheduled {
4619 attempt: 3,
4620 max_retries: 10,
4621 delay_ms: 8_000,
4622 error: "503 service unavailable".into(),
4623 });
4624 assert_eq!(retry["type"], "retry_scheduled");
4625 assert_eq!(retry["attempt"], 3);
4626 assert_eq!(retry["maxRetries"], 10);
4627 assert_eq!(retry["delayMs"], 8_000);
4628 assert_eq!(retry["error"], "503 service unavailable");
4629 }
4630}
4631
4632#[cfg(test)]
4633mod page_window_tests {
4634 use super::page_window;
4635
4636 #[test]
4637 fn tail_is_kept_without_before_id() {
4638 assert_eq!(page_window(12, 5, None), (7..12, true));
4640 assert_eq!(page_window(3, 5, None), (0..3, false));
4642 }
4643
4644 #[test]
4645 fn before_id_opens_the_window_immediately_before_it() {
4646 assert_eq!(page_window(12, 4, Some(10)), (6..10, true));
4648 assert_eq!(page_window(12, 10, Some(10)), (0..10, false));
4650 }
4651
4652 #[test]
4653 fn limit_larger_than_the_branch_never_underflows() {
4654 assert_eq!(page_window(5, 100, Some(2)), (0..2, false));
4655 assert_eq!(page_window(5, 100, None), (0..5, false));
4656 }
4657}
4658
4659#[cfg(test)]
4660mod session_storage_notice_tests {
4661 use super::session_storage_notice;
4662
4663 #[test]
4667 fn the_headless_notice_does_not_promise_an_unavailable_recourse() {
4668 let notice = session_storage_notice("disk full");
4669 assert!(notice.contains("disk full"), "{notice}");
4670 assert!(
4671 !notice.contains("/export"),
4672 "print mode has no interactive /export: {notice}"
4673 );
4674 assert!(
4675 notice.contains("print mode exits"),
4676 "the notice must say the turn is not coming back: {notice}"
4677 );
4678 }
4679}