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