public/galaxy
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Lots of changes... not done yet.

Browse Source
This commit is contained in:
Ryan Ward
2026-08-17 18:19:37 -05:00
parent b5f3290d1a
commit 56e3b51d48
55 changed files with 4494 additions and 1098 deletions
Download Patch File Download Diff File Expand all files Collapse all files
app/src/ai/runtime/event_translator.rs
+102 -9
View File
@@ -62,9 +62,12 @@ impl ProviderRunResponseProjector {
}
}
pub(crate) fn restored(config: RuntimeResponseConfig) -> Self {
pub(crate) fn restored(
config: RuntimeResponseConfig,
projection_was_initialized: bool,
) -> Self {
Self {
translator: RuntimeResponseTranslator::restored(config),
translator: RuntimeResponseTranslator::restored(config, projection_was_initialized),
has_started_model_turn: false,
finished: false,
}
@@ -88,9 +91,11 @@ impl ProviderRunResponseProjector {
})
}
ProviderRunProjection::ModelEvent { event, .. } => self.translator.translate(event),
ProviderRunProjection::ModelRetry { .. } => {
Ok(self.translator.discard_failed_turn_output())
}
ProviderRunProjection::ModelTurnRequested { .. }
| ProviderRunProjection::ModelTurnFinished { .. }
| ProviderRunProjection::ModelRetry { .. }
| ProviderRunProjection::ToolBatchReady { .. } => Ok(Vec::new()),
}
}
@@ -130,8 +135,15 @@ impl RuntimeResponseTranslator {
Self::with_initialization(config, false)
}
pub(crate) fn restored(config: RuntimeResponseConfig) -> Self {
Self::with_initialization(config, true)
pub(crate) fn restored(
mut config: RuntimeResponseConfig,
projection_was_initialized: bool,
) -> Self {
// The task and exchange already exist in restored history. If its output was never
// initialized, replay only the stream Init rather than duplicating task/input messages.
config.needs_create_task = false;
config.user_query = None;
Self::with_initialization(config, projection_was_initialized)
}
fn with_initialization(config: RuntimeResponseConfig, initialized: bool) -> Self {
@@ -143,7 +155,7 @@ impl RuntimeResponseTranslator {
reasoning_message_id: None,
activity_message_ids: HashMap::new(),
activities: HashMap::new(),
has_visible_output: initialized,
has_visible_output: false,
usage: Usage::default(),
context_usage: None,
}
@@ -163,9 +175,7 @@ impl RuntimeResponseTranslator {
}
AgentEvent::ReasoningCompleted { text, .. } => {
self.initialize(&mut events);
if self.reasoning_message_id.is_none() && !text.is_empty() {
self.add_or_append_reasoning(&text, &mut events);
}
self.complete_reasoning(&text, &mut events);
}
AgentEvent::RuntimeActivityUpdated { activity } => {
if self.config.capabilities.host_tool_execution {
@@ -244,6 +254,25 @@ impl RuntimeResponseTranslator {
self.reasoning_message_id = None;
}
fn discard_failed_turn_output(&mut self) -> Vec<ResponseEvent> {
let mut events = Vec::new();
if let Some(message_id) = self.text_message_id.take() {
events.push(build_replace_text_message(
&self.config.task_id,
&message_id,
"",
));
}
if let Some(message_id) = self.reasoning_message_id.take() {
events.push(build_replace_reasoning_message(
&self.config.task_id,
&message_id,
"",
));
}
events
}
pub(crate) fn set_task_id(&mut self, task_id: impl Into<String>) {
let task_id = task_id.into();
if self.config.task_id == task_id {
@@ -315,6 +344,18 @@ impl RuntimeResponseTranslator {
}
}
fn complete_reasoning(&mut self, text: &str, events: &mut Vec<ResponseEvent>) {
if let Some(message_id) = &self.reasoning_message_id {
events.push(build_replace_reasoning_message(
&self.config.task_id,
message_id,
text,
));
} else if !text.is_empty() {
self.add_or_append_reasoning(text, events);
}
}
fn upsert_runtime_activity(
&mut self,
activity: RuntimeActivity,
@@ -441,6 +482,58 @@ fn build_reasoning_message(
runtime_client_action(action)
}
fn build_replace_reasoning_message(task_id: &str, message_id: &str, text: &str) -> ResponseEvent {
let message = api::Message {
id: message_id.to_owned(),
task_id: task_id.to_owned(),
request_id: String::new(),
timestamp: None,
server_message_data: String::new(),
citations: Vec::new(),
fetched_memories: Vec::new(),
message: Some(api::message::Message::AgentReasoning(
api::message::AgentReasoning {
reasoning: text.to_owned(),
finished_duration: None,
},
)),
};
runtime_client_action(api::client_action::Action::UpdateTaskMessage(
api::client_action::UpdateTaskMessage {
task_id: task_id.to_owned(),
message: Some(message),
mask: Some(prost_types::FieldMask {
paths: vec!["agent_reasoning.reasoning".to_owned()],
}),
},
))
}
fn build_replace_text_message(task_id: &str, message_id: &str, text: &str) -> ResponseEvent {
runtime_client_action(api::client_action::Action::UpdateTaskMessage(
api::client_action::UpdateTaskMessage {
task_id: task_id.to_owned(),
message: Some(api::Message {
id: message_id.to_owned(),
task_id: task_id.to_owned(),
request_id: String::new(),
timestamp: None,
server_message_data: String::new(),
citations: Vec::new(),
fetched_memories: Vec::new(),
message: Some(api::message::Message::AgentOutput(
api::message::AgentOutput {
text: text.to_owned(),
},
)),
}),
mask: Some(prost_types::FieldMask {
paths: vec!["agent_output.text".to_owned()],
}),
},
))
}
fn runtime_activity_fallback_text(activity: &RuntimeActivity) -> String {
let title = &activity.title;
let status = activity.status.as_ref().map(|status| match status {
app/src/ai/runtime/event_translator_tests.rs
+172 -1
View File
@@ -45,7 +45,7 @@ fn restored_provider_projection_skips_stream_initialization() {
capabilities: RuntimeCapabilities::provider(),
empty_output_message: None,
};
let mut projector = ProviderRunResponseProjector::restored(config);
let mut projector = ProviderRunResponseProjector::restored(config, true);
let work_id = galaxy_agent_core::ExternalWorkId {
run_id: galaxy_agent_core::ProviderRunId::new("run"),
epoch: galaxy_agent_core::RunEpoch::new(2),
@@ -82,6 +82,56 @@ fn restored_provider_projection_skips_stream_initialization() {
));
}
#[test]
fn restored_uninitialized_projection_replays_init_before_live_delta() {
let config = RuntimeResponseConfig {
task_id: "task".to_owned(),
conversation_id: "conversation".to_owned(),
needs_create_task: true,
user_query: Some("already persisted".to_owned()),
model_id: "model".to_owned(),
max_context_tokens: Some(1_000),
capabilities: RuntimeCapabilities::provider(),
empty_output_message: None,
};
let mut projector = ProviderRunResponseProjector::restored(config, false);
let work_id = galaxy_agent_core::ExternalWorkId {
run_id: galaxy_agent_core::ProviderRunId::new("run"),
epoch: galaxy_agent_core::RunEpoch::new(2),
};
let started = projector
.project(ProviderRunProjection::ModelTurnStarted {
work_id: work_id.clone(),
profile: galaxy_agent_core::ProviderRequestProfile::new("base"),
runtime_id: "runtime".to_owned(),
model_id: "model".to_owned(),
runtime_request_id: "request".to_owned(),
retry_attempt: 0,
elapsed_ms: 1,
})
.unwrap();
let delta = projector
.project(ProviderRunProjection::ModelEvent {
work_id,
event: AgentEvent::TextDelta {
text: "continued".to_owned(),
},
})
.unwrap();
assert_eq!(started.len(), 1);
assert!(matches!(
started[0].r#type,
Some(response_event::Type::Init(_))
));
assert_eq!(delta.len(), 1);
assert!(matches!(
delta[0].r#type,
Some(response_event::Type::ClientActions(_))
));
}
#[test]
fn provider_followup_turn_starts_a_distinct_text_message() {
let mut projector = ProviderRunResponseProjector::new(RuntimeResponseConfig {
@@ -261,6 +311,127 @@ fn reasoning_uses_the_native_reasoning_message_contract() {
));
}
#[test]
fn reasoning_completed_authoritatively_replaces_streamed_reasoning() {
let mut translator = provider_translator();
let streamed = translator
.translate(AgentEvent::ReasoningDelta {
text: "draft reasoning".to_owned(),
})
.expect("reasoning delta");
let completed = translator
.translate(AgentEvent::ReasoningCompleted {
text: "authoritative reasoning".to_owned(),
signature: Some("signature".to_owned()),
})
.expect("reasoning completion");
let Some(response_event::Type::ClientActions(streamed_actions)) = &streamed[1].r#type else {
panic!("expected streamed reasoning action");
};
let Some(client_action::Action::AddMessagesToTask(add)) = &streamed_actions.actions[0].action
else {
panic!("expected streamed reasoning message");
};
let message_id = add.messages[0].id.clone();
let Some(response_event::Type::ClientActions(completed_actions)) = &completed[0].r#type else {
panic!("expected completed reasoning action");
};
let Some(client_action::Action::UpdateTaskMessage(update)) =
&completed_actions.actions[0].action
else {
panic!("completed reasoning must replace the streamed value");
};
let message = update.message.as_ref().expect("replacement message");
assert_eq!(message.id, message_id);
assert!(matches!(
&message.message,
Some(message::Message::AgentReasoning(reasoning))
if reasoning.reasoning == "authoritative reasoning"
));
assert_eq!(
update.mask.as_ref().expect("replacement mask").paths,
["agent_reasoning.reasoning"]
);
}
#[test]
fn provider_retry_clears_failed_attempt_output_before_new_messages() {
let mut projector = ProviderRunResponseProjector::new(RuntimeResponseConfig {
task_id: "task".to_owned(),
conversation_id: "conversation".to_owned(),
needs_create_task: false,
user_query: None,
model_id: "model".to_owned(),
max_context_tokens: Some(1_000),
capabilities: RuntimeCapabilities::provider(),
empty_output_message: None,
});
let work_id = galaxy_agent_core::ExternalWorkId {
run_id: galaxy_agent_core::ProviderRunId::new("run"),
epoch: galaxy_agent_core::RunEpoch::new(1),
};
projector
.project(ProviderRunProjection::ModelTurnStarted {
work_id: work_id.clone(),
profile: galaxy_agent_core::ProviderRequestProfile::new("base"),
runtime_id: "runtime".to_owned(),
model_id: "model".to_owned(),
runtime_request_id: "request".to_owned(),
retry_attempt: 0,
elapsed_ms: 1,
})
.unwrap();
for event in [
AgentEvent::TextDelta {
text: "failed text".to_owned(),
},
AgentEvent::ReasoningDelta {
text: "failed reasoning".to_owned(),
},
] {
projector
.project(ProviderRunProjection::ModelEvent {
work_id: work_id.clone(),
event,
})
.unwrap();
}
let retry = projector
.project(ProviderRunProjection::ModelRetry {
work_id,
profile: galaxy_agent_core::ProviderRequestProfile::new("base"),
runtime_id: "runtime".to_owned(),
model_id: "model".to_owned(),
retry_attempt: 1,
elapsed_ms: 2,
error: galaxy_agent_core::AgentError::new(
galaxy_agent_core::AgentErrorKind::Transport,
"retry",
),
})
.unwrap();
assert_eq!(retry.len(), 2);
for event in retry {
let Some(response_event::Type::ClientActions(actions)) = event.r#type else {
panic!("cleanup must use a client action");
};
let Some(client_action::Action::UpdateTaskMessage(update)) = &actions.actions[0].action
else {
panic!("cleanup must replace failed output");
};
let message = update.message.as_ref().expect("cleanup message");
match &message.message {
Some(message::Message::AgentOutput(output)) => assert!(output.text.is_empty()),
Some(message::Message::AgentReasoning(reasoning)) => {
assert!(reasoning.reasoning.is_empty())
}
_ => panic!("cleanup must target visible text or reasoning"),
}
}
}
#[test]
fn session_activity_updates_the_same_structured_message() {
let mut translator = session_translator();
app/src/ai/runtime/provider_run_coordinator.rs
+276 -135
View File
@@ -5,16 +5,17 @@ use std::sync::Arc;
use std::time::Duration;
use futures::future::BoxFuture;
use futures::StreamExt;
use futures::{FutureExt as _, StreamExt};
use galaxy_agent_core::{
AgentError, AgentErrorKind, AgentEvent, AgentRuntime, CompletedModelTurn, ContentPart,
ExternalWorkId, ModelFailureDisposition, PendingToolBatch, ProviderModelCall,
turn_control, AgentError, AgentErrorKind, AgentEvent, AgentRuntime, CompletedModelTurn,
ContentPart, ExternalWorkId, ModelFailureDisposition, PendingToolBatch, ProviderModelCall,
ProviderRequestProfile, ProviderRun, ProviderRunFailureKind, ProviderRunId, ProviderRunLimits,
ProviderRunOutcome, ProviderRunPhase, ProviderRunProtocolError, ProviderRunState,
ProviderRunStep, RunEpoch, RuntimeKind, StopReason, ToolEvent, TurnControl, TurnRequest, Usage,
ProviderRunStep, RunEpoch, RuntimeKind, StopReason, ToolEvent, TurnCommand, TurnControl,
TurnRequest, Usage,
};
use instant::Instant;
use warpui::r#async::FutureExt as _;
use warpui::r#async::Timer;
use crate::ai::agent::conversation::AIConversationId;
@@ -321,31 +322,55 @@ impl ProviderRunCoordinator {
&mut self,
control: TurnControl,
mut project: F,
mut checkpoint: C,
checkpoint: C,
) -> Result<ProviderRunBlock, ProviderRunCoordinatorError>
where
F: FnMut(ProviderRunProjection) -> Result<(), String>,
C: FnMut(ProviderRun) -> BoxFuture<'static, Result<(), String>>,
{
self.drive_until_blocked_with_acknowledgements(
control,
move |projection| {
let result = project(projection);
Box::pin(async move { result })
},
checkpoint,
)
.await
}
pub(crate) async fn drive_until_blocked_with_acknowledgements<F, C>(
&mut self,
control: TurnControl,
mut project: F,
mut checkpoint: C,
) -> Result<ProviderRunBlock, ProviderRunCoordinatorError>
where
F: FnMut(ProviderRunProjection) -> BoxFuture<'static, Result<(), String>>,
C: FnMut(ProviderRun) -> BoxFuture<'static, Result<(), String>>,
{
loop {
match self.run.next_step()? {
Some(ProviderRunStep::CallModel(call)) => {
if !self.checkpoint_or_fail(&mut checkpoint).await? {
continue;
return self.terminal_block();
}
self.drive_model_call(call, control.clone(), &mut project)
self.drive_model_call_acknowledged(call, control.clone(), &mut project)
.await?;
}
Some(ProviderRunStep::DispatchTools(batch)) => {
if !self.checkpoint_or_fail(&mut checkpoint).await? {
continue;
return self.terminal_block();
}
if !self.project_or_fail(
ProviderRunProjection::ToolBatchReady {
batch: batch.clone(),
},
&mut project,
)? {
if !self
.project_or_fail_acknowledged(
ProviderRunProjection::ToolBatchReady {
batch: batch.clone(),
},
&mut project,
)
.await?
{
continue;
}
if batch.is_complete() {
@@ -356,7 +381,7 @@ impl ProviderRunCoordinator {
}
Some(ProviderRunStep::Done(outcome)) => {
if !self.checkpoint_or_fail(&mut checkpoint).await? {
continue;
return self.terminal_block();
}
return Ok(ProviderRunBlock::Done(outcome));
}
@@ -368,7 +393,7 @@ impl ProviderRunCoordinator {
let work_id = work_id.clone();
let stop_reason = stop_reason.clone();
if !self.checkpoint_or_fail(&mut checkpoint).await? {
continue;
return self.terminal_block();
}
return Ok(ProviderRunBlock::AwaitingDriver {
work_id,
@@ -413,14 +438,26 @@ impl ProviderRunCoordinator {
}
}
async fn drive_model_call<F>(
fn terminal_block(&mut self) -> Result<ProviderRunBlock, ProviderRunCoordinatorError> {
match self.run.next_step()? {
Some(ProviderRunStep::Done(outcome)) => Ok(ProviderRunBlock::Done(outcome)),
Some(ProviderRunStep::CallModel(_) | ProviderRunStep::DispatchTools(_)) | None => Err(
ProviderRunCoordinatorError::Core(ProviderRunProtocolError::UnexpectedState {
expected: ProviderRunPhase::Failed,
actual: self.run.state().phase(),
}),
),
}
}
async fn drive_model_call_acknowledged<F>(
&mut self,
call: ProviderModelCall,
control: TurnControl,
project: &mut F,
) -> Result<(), ProviderRunCoordinatorError>
where
F: FnMut(ProviderRunProjection) -> Result<(), String>,
F: FnMut(ProviderRunProjection) -> BoxFuture<'static, Result<(), String>>,
{
let Some(profile) = self.profiles.get(call.profile.as_str()).cloned() else {
self.run.fail(
@@ -440,39 +477,66 @@ impl ProviderRunCoordinator {
.collect::<BTreeSet<_>>();
let runtime_id = profile.runtime.descriptor().id.clone();
let model_id = profile.request.model.as_str().to_string();
if !self.project_or_fail(
ProviderRunProjection::ModelTurnRequested {
work_id: call.work_id.clone(),
profile: call.profile.clone(),
runtime_id: runtime_id.clone(),
model_id: model_id.clone(),
retry_attempt: call.retry_attempt,
},
project,
)? {
if !self
.project_or_fail_acknowledged(
ProviderRunProjection::ModelTurnRequested {
work_id: call.work_id.clone(),
profile: call.profile.clone(),
runtime_id: runtime_id.clone(),
model_id: model_id.clone(),
retry_attempt: call.retry_attempt,
},
project,
)
.await?
{
return Ok(());
}
let request = request_for_model_call(profile.request.clone(), &call);
let started_at = Instant::now();
let stream = match profile
.runtime
.start_turn(request, control)
.with_timeout(self.model_start_timeout)
.await
{
Ok(Ok(stream)) => stream,
Ok(Err(error)) => {
self.handle_model_failure(&call, &profile, started_at, error, project)?;
let (attempt_sender, attempt_control) = turn_control();
let start_future = profile.runtime.start_turn(request, attempt_control).fuse();
let timeout = futures::FutureExt::fuse(Timer::after(self.model_start_timeout));
futures::pin_mut!(start_future, timeout);
let mut control_open = true;
let start_result = loop {
let command = if control_open {
futures::future::Either::Left(control.receive())
} else {
futures::future::Either::Right(futures::future::pending())
}
.fuse();
futures::pin_mut!(command);
futures::select_biased! {
result = start_future => break Some(result),
_ = timeout => {
let _ = attempt_sender.send(TurnCommand::Cancel).await;
break None;
}
command = command => match command {
Ok(command) => {
let _ = attempt_sender.send(command).await;
}
Err(_) => control_open = false,
},
}
};
let stream = match start_result {
Some(Ok(stream)) => stream,
Some(Err(error)) => {
self.handle_model_failure_acknowledged(&call, &profile, started_at, error, project)
.await?;
return Ok(());
}
Err(_) => {
self.handle_model_failure(
None => {
self.handle_model_failure_acknowledged(
&call,
&profile,
started_at,
provider_timeout_error("start", self.model_start_timeout),
project,
)?;
)
.await?;
return Ok(());
}
};
@@ -480,157 +544,230 @@ impl ProviderRunCoordinator {
let mut buffer = ModelTurnBuffer::default();
loop {
let event = match stream
.next()
.with_timeout(self.model_event_idle_timeout)
.await
{
Ok(Some(Ok(event))) => event,
Ok(Some(Err(error))) => {
self.handle_model_failure(&call, &profile, started_at, error, project)?;
let next_event = stream.next().fuse();
let timeout = futures::FutureExt::fuse(Timer::after(self.model_event_idle_timeout));
futures::pin_mut!(next_event, timeout);
let event_result = loop {
let command = if control_open {
futures::future::Either::Left(control.receive())
} else {
futures::future::Either::Right(futures::future::pending())
}
.fuse();
futures::pin_mut!(command);
futures::select_biased! {
event = next_event => break Some(event),
_ = timeout => {
let _ = attempt_sender.send(TurnCommand::Cancel).await;
break None;
}
command = command => match command {
Ok(command) => {
let _ = attempt_sender.send(command).await;
}
Err(_) => control_open = false,
},
}
};
let event = match event_result {
Some(Some(Ok(event))) => event,
Some(Some(Err(error))) => {
self.handle_model_failure_acknowledged(
&call, &profile, started_at, error, project,
)
.await?;
return Ok(());
}
Ok(None) => break,
Err(_) => {
self.handle_model_failure(
Some(None) => break,
None => {
self.handle_model_failure_acknowledged(
&call,
&profile,
started_at,
provider_timeout_error("event", self.model_event_idle_timeout),
project,
)?;
)
.await?;
return Ok(());
}
};
match event {
AgentEvent::TurnStarted { runtime_request_id } => {
if buffer.started {
self.handle_model_failure(
self.handle_model_failure_acknowledged(
&call,
&profile,
started_at,
protocol_error("provider emitted more than one TurnStarted event"),
project,
)?;
)
.await?;
return Ok(());
}
if runtime_request_id.is_empty() {
self.handle_model_failure(
self.handle_model_failure_acknowledged(
&call,
&profile,
started_at,
protocol_error("provider emitted an empty runtime request ID"),
project,
)?;
)
.await?;
return Ok(());
}
buffer.started = true;
if !self.project_or_fail(
ProviderRunProjection::ModelTurnStarted {
work_id: call.work_id.clone(),
profile: call.profile.clone(),
runtime_id: runtime_id.clone(),
model_id: model_id.clone(),
runtime_request_id,
retry_attempt: call.retry_attempt,
elapsed_ms: elapsed_millis(started_at),
},
project,
)? {
if !self
.project_or_fail_acknowledged(
ProviderRunProjection::ModelTurnStarted {
work_id: call.work_id.clone(),
profile: call.profile.clone(),
runtime_id: runtime_id.clone(),
model_id: model_id.clone(),
runtime_request_id,
retry_attempt: call.retry_attempt,
elapsed_ms: elapsed_millis(started_at),
},
project,
)
.await?
{
return Ok(());
}
}
AgentEvent::TextDelta { text } => {
if !self.ensure_model_started(&call, &profile, started_at, &buffer, project)? {
if !self
.ensure_model_started_acknowledged(
&call, &profile, started_at, &buffer, project,
)
.await?
{
return Ok(());
}
buffer.text.push_str(&text);
if !self.project_or_fail(
ProviderRunProjection::ModelEvent {
work_id: call.work_id.clone(),
event: AgentEvent::TextDelta { text },
},
project,
)? {
if !self
.project_or_fail_acknowledged(
ProviderRunProjection::ModelEvent {
work_id: call.work_id.clone(),
event: AgentEvent::TextDelta { text },
},
project,
)
.await?
{
return Ok(());
}
}
AgentEvent::ReasoningDelta { text } => {
if !self.ensure_model_started(&call, &profile, started_at, &buffer, project)? {
if !self
.ensure_model_started_acknowledged(
&call, &profile, started_at, &buffer, project,
)
.await?
{
return Ok(());
}
buffer.reasoning.push_str(&text);
if !self.project_or_fail(
ProviderRunProjection::ModelEvent {
work_id: call.work_id.clone(),
event: AgentEvent::ReasoningDelta { text },
},
project,
)? {
if !self
.project_or_fail_acknowledged(
ProviderRunProjection::ModelEvent {
work_id: call.work_id.clone(),
event: AgentEvent::ReasoningDelta { text },
},
project,
)
.await?
{
return Ok(());
}
}
AgentEvent::ReasoningCompleted { text, signature } => {
if !self.ensure_model_started(&call, &profile, started_at, &buffer, project)? {
if !self
.ensure_model_started_acknowledged(
&call, &profile, started_at, &buffer, project,
)
.await?
{
return Ok(());
}
if !text.is_empty() {
buffer.reasoning.clone_from(&text);
}
buffer.reasoning.clone_from(&text);
buffer.reasoning_signature.clone_from(&signature);
if !self.project_or_fail(
ProviderRunProjection::ModelEvent {
work_id: call.work_id.clone(),
event: AgentEvent::ReasoningCompleted { text, signature },
},
project,
)? {
if !self
.project_or_fail_acknowledged(
ProviderRunProjection::ModelEvent {
work_id: call.work_id.clone(),
event: AgentEvent::ReasoningCompleted { text, signature },
},
project,
)
.await?
{
return Ok(());
}
}
AgentEvent::Tool {
event: ToolEvent::Proposed { call: tool_call },
} => {
if !self.ensure_model_started(&call, &profile, started_at, &buffer, project)? {
if !self
.ensure_model_started_acknowledged(
&call, &profile, started_at, &buffer, project,
)
.await?
{
return Ok(());
}
buffer.tool_calls.push(tool_call);
}
AgentEvent::UsageUpdated { usage } => {
if !self.ensure_model_started(&call, &profile, started_at, &buffer, project)? {
if !self
.ensure_model_started_acknowledged(
&call, &profile, started_at, &buffer, project,
)
.await?
{
return Ok(());
}
buffer.usage.clone_from(&usage);
let cumulative_usage = combined_usage(self.run.usage(), &usage);
if !self.project_or_fail(
ProviderRunProjection::ModelEvent {
work_id: call.work_id.clone(),
event: AgentEvent::UsageUpdated {
usage: cumulative_usage,
if !self
.project_or_fail_acknowledged(
ProviderRunProjection::ModelEvent {
work_id: call.work_id.clone(),
event: AgentEvent::UsageUpdated {
usage: cumulative_usage,
},
},
},
project,
)? {
project,
)
.await?
{
return Ok(());
}
}
AgentEvent::TurnStopped { reason } => {
if !self.ensure_model_started(&call, &profile, started_at, &buffer, project)? {
if !self
.ensure_model_started_acknowledged(
&call, &profile, started_at, &buffer, project,
)
.await?
{
return Ok(());
}
if !self.project_or_fail(
ProviderRunProjection::ModelTurnFinished {
work_id: call.work_id.clone(),
profile: call.profile.clone(),
runtime_id: runtime_id.clone(),
model_id: model_id.clone(),
stop_reason: reason.clone(),
retry_attempt: call.retry_attempt,
elapsed_ms: elapsed_millis(started_at),
tool_call_count: buffer.tool_calls.len(),
},
project,
)? {
if !self
.project_or_fail_acknowledged(
ProviderRunProjection::ModelTurnFinished {
work_id: call.work_id.clone(),
profile: call.profile.clone(),
runtime_id: runtime_id.clone(),
model_id: model_id.clone(),
stop_reason: reason.clone(),
retry_attempt: call.retry_attempt,
elapsed_ms: elapsed_millis(started_at),
tool_call_count: buffer.tool_calls.len(),
},
project,
)
.await?
{
return Ok(());
}
if reason == StopReason::Cancelled {
@@ -657,7 +794,7 @@ impl ProviderRunCoordinator {
| AgentEvent::ContextUsageUpdated { .. }
| AgentEvent::UserInputAccepted { .. }
| AgentEvent::RuntimeNotice { .. } => {
self.handle_model_failure(
self.handle_model_failure_acknowledged(
&call,
&profile,
started_at,
@@ -665,23 +802,25 @@ impl ProviderRunCoordinator {
"direct-provider transport emitted a non-model lifecycle event",
),
project,
)?;
)
.await?;
return Ok(());
}
}
}
self.handle_model_failure(
self.handle_model_failure_acknowledged(
&call,
&profile,
started_at,
protocol_error("provider stream ended before TurnStopped"),
project,
)?;
)
.await?;
Ok(())
}
fn ensure_model_started<F>(
async fn ensure_model_started_acknowledged<F>(
&mut self,
call: &ProviderModelCall,
profile: &ProviderRunProfile,
@@ -690,22 +829,23 @@ impl ProviderRunCoordinator {
project: &mut F,
) -> Result<bool, ProviderRunCoordinatorError>
where
F: FnMut(ProviderRunProjection) -> Result<(), String>,
F: FnMut(ProviderRunProjection) -> BoxFuture<'static, Result<(), String>>,
{
if buffer.started {
return Ok(true);
}
self.handle_model_failure(
self.handle_model_failure_acknowledged(
call,
profile,
started_at,
protocol_error("provider emitted model output before TurnStarted"),
project,
)?;
)
.await?;
Ok(false)
}
fn handle_model_failure<F>(
async fn handle_model_failure_acknowledged<F>(
&mut self,
call: &ProviderModelCall,
profile: &ProviderRunProfile,
@@ -714,7 +854,7 @@ impl ProviderRunCoordinator {
project: &mut F,
) -> Result<(), ProviderRunCoordinatorError>
where
F: FnMut(ProviderRunProjection) -> Result<(), String>,
F: FnMut(ProviderRunProjection) -> BoxFuture<'static, Result<(), String>>,
{
let disposition = self
.run
@@ -737,7 +877,7 @@ impl ProviderRunCoordinator {
));
}
};
self.project_or_fail(
self.project_or_fail_acknowledged(
ProviderRunProjection::ModelRetry {
work_id: call.work_id.clone(),
profile: call.profile.clone(),
@@ -748,20 +888,21 @@ impl ProviderRunCoordinator {
error,
},
project,
)?;
)
.await?;
}
Ok(())
}
fn project_or_fail<F>(
async fn project_or_fail_acknowledged<F>(
&mut self,
event: ProviderRunProjection,
project: &mut F,
) -> Result<bool, ProviderRunCoordinatorError>
where
F: FnMut(ProviderRunProjection) -> Result<(), String>,
F: FnMut(ProviderRunProjection) -> BoxFuture<'static, Result<(), String>>,
{
match project(event) {
match project(event).await {
Ok(()) => Ok(true),
Err(message) => {
self.run.fail(
app/src/ai/runtime/provider_run_coordinator_tests.rs
+140 -1
View File
@@ -79,6 +79,7 @@ struct StallingRuntime {
first_attempt_stall: FirstAttemptStall,
attempts: AtomicUsize,
requests: Mutex<Vec<TurnRequest>>,
controls: Mutex<Vec<TurnControl>>,
}
impl StallingRuntime {
@@ -93,12 +94,22 @@ impl StallingRuntime {
first_attempt_stall,
attempts: AtomicUsize::new(0),
requests: Mutex::new(Vec::new()),
controls: Mutex::new(Vec::new()),
}
}
fn requests(&self) -> Vec<TurnRequest> {
self.requests.lock().unwrap().clone()
}
fn attempt_commands(&self) -> Vec<Option<TurnCommand>> {
self.controls
.lock()
.unwrap()
.iter()
.map(|control| control.try_receive().ok())
.collect()
}
}
#[async_trait]
@@ -110,9 +121,10 @@ impl AgentRuntime for StallingRuntime {
async fn start_turn(
&self,
request: TurnRequest,
_control: TurnControl,
control: TurnControl,
) -> Result<AgentEventStream, AgentError> {
self.requests.lock().unwrap().push(request);
self.controls.lock().unwrap().push(control);
let attempt = self.attempts.fetch_add(1, Ordering::Relaxed);
if attempt == 0 {
match self.first_attempt_stall {
@@ -778,6 +790,64 @@ async fn recoverable_start_failure_retries_the_same_work_identity() {
assert_eq!(retry.1, 1);
}
#[tokio::test]
async fn model_progress_waits_for_each_projection_acknowledgement() {
let runtime = Arc::new(ScriptedRuntime::new(vec![answer_turn()]));
let mut coordinator = coordinator(runtime.clone());
let (acknowledgement_sender, acknowledgement_receiver) = async_channel::unbounded();
let (observed_sender, observed_receiver) = async_channel::unbounded();
let (_sender, control) = turn_control();
let drive = async {
coordinator
.drive_until_blocked_with_acknowledgements(
control,
move |projection| {
let acknowledgement_receiver = acknowledgement_receiver.clone();
let observed_sender = observed_sender.clone();
Box::pin(async move {
observed_sender.send(projection).await.unwrap();
acknowledgement_receiver.recv().await.unwrap()
})
},
|_| Box::pin(async { Ok(()) }),
)
.await
.unwrap()
};
let driver = async {
let first = observed_receiver.recv().await.unwrap();
assert!(matches!(
first,
ProviderRunProjection::ModelTurnRequested { .. }
));
assert_eq!(runtime.requests().len(), 0);
assert!(observed_receiver.try_recv().is_err());
acknowledgement_sender.send(Ok(())).await.unwrap();
let second = observed_receiver.recv().await.unwrap();
assert!(matches!(
second,
ProviderRunProjection::ModelTurnStarted { .. }
));
assert_eq!(runtime.requests().len(), 1);
assert!(observed_receiver.try_recv().is_err());
acknowledgement_sender.send(Ok(())).await.unwrap();
loop {
let projection = observed_receiver.recv().await.unwrap();
let finished = matches!(projection, ProviderRunProjection::ModelTurnFinished { .. });
acknowledgement_sender.send(Ok(())).await.unwrap();
if finished {
break;
}
}
};
let (block, ()) = futures::join!(drive, driver);
assert!(matches!(block, ProviderRunBlock::AwaitingDriver { .. }));
}
fn assert_single_retry_lifecycle(
projections: &[ProviderRunProjection],
expected_timeout_stage: &str,
@@ -868,6 +938,10 @@ async fn model_start_timeout_retries_the_same_work_identity() {
assert!(matches!(block, ProviderRunBlock::AwaitingDriver { .. }));
assert_eq!(runtime.requests().len(), 2);
assert_eq!(
runtime.attempt_commands(),
vec![Some(TurnCommand::Cancel), None]
);
assert_eq!(coordinator.run().model_retries(), 1);
assert_single_retry_lifecycle(&projections, "start timed out", false);
}
@@ -887,10 +961,43 @@ async fn model_event_idle_timeout_retries_the_same_work_identity() {
assert!(matches!(block, ProviderRunBlock::AwaitingDriver { .. }));
assert_eq!(runtime.requests().len(), 2);
assert_eq!(
runtime.attempt_commands(),
vec![Some(TurnCommand::Cancel), None]
);
assert_eq!(coordinator.run().model_retries(), 1);
assert_single_retry_lifecycle(&projections, "event timed out", true);
}
#[tokio::test]
async fn persistent_checkpoint_failure_terminates_without_redrive() {
let runtime = Arc::new(ScriptedRuntime::new(vec![answer_turn()]));
let mut coordinator = coordinator(runtime.clone());
let checkpoint_attempts = Arc::new(AtomicUsize::new(0));
let attempts = checkpoint_attempts.clone();
let (_sender, control) = turn_control();
let block = coordinator
.drive_until_blocked_with_checkpoint(
control,
|_| Ok(()),
move |_| {
attempts.fetch_add(1, Ordering::Relaxed);
Box::pin(async { Err("database unavailable".to_string()) })
},
)
.await
.unwrap();
let ProviderRunBlock::Done(ProviderRunOutcome::Failed(failure)) = block else {
panic!("checkpoint failure must terminate the drive");
};
assert_eq!(failure.kind, ProviderRunFailureKind::ExternalWork);
assert!(failure.message.contains("database unavailable"));
assert_eq!(checkpoint_attempts.load(Ordering::Relaxed), 1);
assert!(runtime.requests().is_empty());
}
#[tokio::test]
async fn tool_projection_failure_synthesizes_a_result_and_fails_the_run() {
let runtime = Arc::new(ScriptedRuntime::new(vec![tool_turn()]));
@@ -953,6 +1060,38 @@ async fn stream_without_terminal_event_fails_instead_of_committing_partial_outpu
assert_eq!(coordinator.run().transcript().len(), 1);
}
#[tokio::test]
async fn reasoning_completed_authoritatively_replaces_canonical_reasoning() {
let runtime = Arc::new(ScriptedRuntime::new(vec![Ok(vec![
started("request-reasoning"),
Ok(AgentEvent::ReasoningDelta {
text: "draft".to_string(),
}),
Ok(AgentEvent::ReasoningCompleted {
text: "final".to_string(),
signature: Some("signature".to_string()),
}),
stopped(StopReason::Completed),
])]));
let mut coordinator = coordinator(runtime);
let (_sender, control) = turn_control();
let block = coordinator
.drive_until_blocked(control, |_| Ok(()))
.await
.unwrap();
assert!(matches!(block, ProviderRunBlock::AwaitingDriver { .. }));
let MessageContent::MultiPart(parts) = &coordinator.run().transcript()[1].content else {
panic!("expected canonical reasoning content");
};
assert!(matches!(
parts.as_slice(),
[ContentPart::Reasoning { text, signature }]
if text == "final" && signature.as_deref() == Some("signature")
));
}
#[tokio::test]
async fn provider_cancellation_does_not_commit_partial_assistant_content() {
let expected_transcript = request().messages;
app/src/ai/runtime/rig_request.rs
+10
View File
@@ -835,6 +835,16 @@ fn build_system_prompt(
.join(", "),
);
prompt.push_str(".\nNever invent tool names or parameters. Check command exit codes and tool error results before claiming success.\n");
if tools.iter().any(|tool| tool.name == "run_shell_command") {
let has_file_tools = tools
.iter()
.any(|tool| matches!(tool.name.as_str(), "file_glob" | "grep" | "read_files"));
if has_file_tools {
prompt.push_str(
"Prefer `file_glob`, `grep`, and `read_files` for file discovery, content search, and file reading when they are available. Reserve `run_shell_command` for operations those specialized tools cannot perform; do not use shell `find`, `grep`, `rg`, `cat`, `head`, or `tail` as substitutes.\n",
);
}
}
if tools.iter().any(|tool| tool.name == "create_plan") {
prompt.push_str(
"Plan document creation is available through `create_plan`. When the user asks you to create a plan for review, research first as needed, then call `create_plan`; do not merely return the plan as prose or claim that no plan-creation tool is available. If the user asks to review the plan before implementation, creating the document and presenting it for review is the requested outcome; do not implement it until they approve.\n",
app/src/ai/runtime/rig_request_tests.rs
+22
View File
@@ -127,6 +127,28 @@ fn builds_a_rig_turn_directly_from_galaxy_request_state() {
));
}
#[test]
fn system_prompt_prefers_specialized_file_tools_over_shell_substitutes() {
let mut params = RequestParams::new_for_test();
params.input = vec![user_query("Find every Rust file containing ProviderRun")];
let prepared = prepare_rig_turn(
&config(),
params,
vec![
ToolType::RunShellCommand,
ToolType::FileGlob,
ToolType::Grep,
ToolType::ReadFiles,
],
Vec::new(),
);
let prompt = prepared.request.system_prompt.expect("system prompt");
assert!(prompt.contains("Prefer `file_glob`, `grep`, and `read_files`"));
assert!(prompt.contains("do not use shell `find`, `grep`, `rg`, `cat`, `head`, or `tail`"));
}
#[test]
fn normal_turn_advertises_plan_creation_and_corrects_false_unavailability_claims() {
let mut params = RequestParams::new_for_test();
app/src/ai/runtime/rig_tool.rs
+479 -299
View File
@@ -25,6 +25,9 @@ pub(crate) fn action_from_tool_call(
mcp_tool_aliases: &HashMap<String, MCPToolTarget>,
) -> Result<AIAgentAction, String> {
let input = &call.arguments;
if !input.is_object() {
return Err(format!("invalid {} input: expected an object", call.name));
}
let action = if let Some(target) = mcp_tool_aliases.get(&call.name) {
AIAgentActionType::CallMCPTool {
server_id: target.server_id,
@@ -33,197 +36,220 @@ pub(crate) fn action_from_tool_call(
}
} else {
match call.name.as_str() {
"run_shell_command" => AIAgentActionType::RequestCommandOutput {
command: string(input, "command"),
is_read_only: Some(boolean(input, "is_read_only")),
is_risky: Some(boolean(input, "is_risky")),
wait_until_completion: boolean(input, "wait_until_complete"),
uses_pager: Some(boolean(input, "uses_pager")),
rationale: None,
citations: Vec::new(),
},
"read_files" => AIAgentActionType::ReadFiles(ReadFilesRequest {
locations: input
.get("files")
.and_then(serde_json::Value::as_array)
.into_iter()
.flatten()
.filter_map(file_location)
.collect(),
}),
"apply_file_diffs" => AIAgentActionType::RequestFileEdits {
file_edits: file_edits(input),
title: nonempty_string(input, "summary"),
},
"grep" => AIAgentActionType::Grep {
queries: strings(input, "queries"),
path: string(input, "path"),
},
"file_glob" => AIAgentActionType::FileGlob {
patterns: strings(input, "patterns"),
path: nonempty_string(input, "path"),
},
"search_codebase" => AIAgentActionType::SearchCodebase(SearchCodebaseRequest {
query: string(input, "query"),
partial_paths: nonempty_strings(input, "path_filters"),
codebase_path: nonempty_string(input, "path"),
}),
"write_to_long_running_shell_command" => {
AIAgentActionType::WriteToLongRunningShellCommand {
block_id: string(input, "command_id").into(),
input: string(input, "input").into_bytes().into(),
mode: match input.get("mode").and_then(serde_json::Value::as_str) {
Some("line") => AIAgentPtyWriteMode::Line,
Some("block") => AIAgentPtyWriteMode::Block,
Some("raw") | Some(_) | None => AIAgentPtyWriteMode::Raw,
},
"run_shell_command" => AIAgentActionType::RequestCommandOutput {
command: required_nonempty_string(input, "command")?,
is_read_only: Some(optional_boolean(input, "is_read_only")?.unwrap_or(false)),
is_risky: Some(optional_boolean(input, "is_risky")?.unwrap_or(false)),
wait_until_completion: optional_boolean(input, "wait_until_complete")?
.unwrap_or(false),
uses_pager: Some(optional_boolean(input, "uses_pager")?.unwrap_or(false)),
rationale: None,
citations: Vec::new(),
},
"read_files" => AIAgentActionType::ReadFiles(ReadFilesRequest {
locations: required_array(input, "files")?
.iter()
.enumerate()
.map(|(index, file)| file_location(file, index))
.collect::<Result<_, _>>()?,
}),
"apply_file_diffs" => AIAgentActionType::RequestFileEdits {
file_edits: file_edits(input)?,
title: Some(required_string(input, "summary")?),
},
"grep" => AIAgentActionType::Grep {
queries: required_strings(input, "queries")?,
path: optional_string(input, "path")?.unwrap_or_default(),
},
"file_glob" => AIAgentActionType::FileGlob {
patterns: required_strings(input, "patterns")?,
path: optional_string(input, "path")?.filter(|path| !path.is_empty()),
},
"search_codebase" => AIAgentActionType::SearchCodebase(SearchCodebaseRequest {
query: required_string(input, "query")?,
partial_paths: optional_strings(input, "path_filters")?
.filter(|paths| !paths.is_empty()),
codebase_path: optional_string(input, "path")?.filter(|path| !path.is_empty()),
}),
"write_to_long_running_shell_command" => {
AIAgentActionType::WriteToLongRunningShellCommand {
block_id: required_nonempty_string(input, "command_id")?.into(),
input: required_string(input, "input")?.into_bytes().into(),
mode: match optional_string(input, "mode")?.as_deref() {
Some("line") => AIAgentPtyWriteMode::Line,
Some("block") => AIAgentPtyWriteMode::Block,
Some("raw") | None => AIAgentPtyWriteMode::Raw,
Some(mode) => {
return Err(format!(
"invalid field \"mode\": expected \"raw\", \"line\", or \"block\", got {mode:?}"
));
}
},
}
}
}
"interrupt_shell_command" => AIAgentActionType::WriteToLongRunningShellCommand {
block_id: string(input, "command_id").into(),
input: vec![galaxy_terminal::model::escape_sequences::C0::ETX].into(),
mode: AIAgentPtyWriteMode::Raw,
},
"read_shell_command_output" => AIAgentActionType::ReadShellCommandOutput {
block_id: string(input, "command_id").into(),
delay: Some(ShellCommandDelay::Duration(Duration::from_secs(
input
.get("wait_seconds")
.and_then(serde_json::Value::as_u64)
.unwrap_or(2)
.min(crate::ai::bedrock::request_translator::COMMAND_MONITOR_MAX_POLL_SECONDS),
))),
},
"read_mcp_resource" => AIAgentActionType::ReadMCPResource {
server_id: uuid(input, "server_id"),
name: String::new(),
uri: nonempty_string(input, "uri"),
},
"read_plan" | "read_documents" | "read_notebook" => {
AIAgentActionType::ReadDocuments(ReadDocumentsRequest {
document_ids: strings(input, "document_ids")
.into_iter()
.filter_map(|id| AIDocumentId::try_from(id).ok())
.collect(),
})
}
"create_plan" | "create_documents" | "create_notebook" => {
AIAgentActionType::CreateDocuments(CreateDocumentsRequest {
documents: input
.get("documents")
.and_then(serde_json::Value::as_array)
.into_iter()
.flatten()
.filter_map(|document| {
Some(DocumentToCreate {
title: document.get("title")?.as_str()?.to_string(),
content: document.get("content")?.as_str()?.to_string(),
})
})
.collect(),
})
}
"edit_plan" | "edit_documents" | "edit_notebook" => {
AIAgentActionType::EditDocuments(EditDocumentsRequest {
diffs: input
.get("diffs")
.and_then(serde_json::Value::as_array)
.into_iter()
.flatten()
.filter_map(|diff| {
Some(DocumentDiff {
document_id: AIDocumentId::try_from(diff.get("document_id")?.as_str()?)
.ok()?,
search: string(diff, "search"),
replace: string(diff, "replace"),
})
})
.collect(),
})
}
"run_agents" => AIAgentActionType::RunAgents(RunAgentsRequest {
summary: string(input, "summary"),
base_prompt: string(input, "base_prompt"),
skills: skill_references(input, skill_path_origin),
model_id: string(input, "model_id"),
harness_type: string(input, "harness_type"),
execution_mode: run_agents_execution_mode(input),
agent_run_configs: input
.get("agent_run_configs")
.and_then(serde_json::Value::as_array)
.into_iter()
.flatten()
.map(|config| RunAgentsAgentRunConfig {
name: string(config, "name"),
prompt: string(config, "prompt"),
title: string(config, "title"),
})
.collect(),
plan_id: string(input, "plan_id"),
harness_auth_secret_name: None,
}),
"start_agent" => AIAgentActionType::StartAgent {
version: StartAgentVersion::V1,
name: string(input, "name"),
prompt: string(input, "prompt"),
execution_mode: StartAgentExecutionMode::local_with_defaults(),
lifecycle_subscription: None,
},
"send_message_to_agent" => AIAgentActionType::SendMessageToAgent {
addresses: vec![string(input, "agent_id")],
subject: String::new(),
message: string(input, "message"),
},
"ask_user_question" => AIAgentActionType::AskUserQuestion {
questions: vec![AskUserQuestionItem {
question_id: Uuid::new_v4().to_string(),
question: string(input, "question"),
question_type: AskUserQuestionType::MultipleChoice {
is_multiselect: false,
options: strings(input, "options")
"interrupt_shell_command" => AIAgentActionType::WriteToLongRunningShellCommand {
block_id: required_nonempty_string(input, "command_id")?.into(),
input: vec![galaxy_terminal::model::escape_sequences::C0::ETX].into(),
mode: AIAgentPtyWriteMode::Raw,
},
"read_shell_command_output" => AIAgentActionType::ReadShellCommandOutput {
block_id: required_nonempty_string(input, "command_id")?.into(),
delay: Some(ShellCommandDelay::Duration(Duration::from_secs(
optional_bounded_u64(
input,
"wait_seconds",
crate::ai::bedrock::request_translator::COMMAND_MONITOR_MAX_POLL_SECONDS,
)?
.unwrap_or(2),
))),
},
"read_mcp_resource" => AIAgentActionType::ReadMCPResource {
server_id: Some(required_uuid(input, "server_id")?),
name: String::new(),
uri: Some(required_string(input, "uri")?),
},
"read_plan" | "read_documents" | "read_notebook" => {
AIAgentActionType::ReadDocuments(ReadDocumentsRequest {
document_ids: required_strings(input, "document_ids")?
.into_iter()
.enumerate()
.map(|(index, label)| AskUserQuestionOption {
label,
recommended: index == 0,
.map(|id| {
AIDocumentId::try_from(id.clone()).map_err(|_| {
format!("invalid document_ids entry: {id:?} is not a document ID")
})
})
.collect(),
supports_other: true,
},
}],
},
"read_skill" => {
let skill = string(input, "skill");
let skill = match input
.get("reference_type")
.and_then(serde_json::Value::as_str)
{
Some("bundled") => SkillReference::BundledSkillId(skill),
Some("path") | Some(_) | None => SkillReference::Path(
skill_path_origin
.location_for_path(skill)
.map_err(|error| error.to_string())?,
),
};
AIAgentActionType::ReadSkill(ReadSkillRequest { skill })
}
"fetch_conversation" => AIAgentActionType::FetchConversation {
conversation_id: string(input, "conversation_id"),
},
name if name.starts_with("mcp__") => {
let mut parts = name.splitn(3, "__");
let _prefix = parts.next();
let server_id = parts.next().and_then(|value| Uuid::parse_str(value).ok());
let name = parts
.next()
.unwrap_or_else(|| name.strip_prefix("mcp__").unwrap_or(name))
.to_string();
AIAgentActionType::CallMCPTool {
server_id,
name,
input: input.clone(),
.collect::<Result<_, _>>()?,
})
}
"create_plan" | "create_documents" | "create_notebook" => {
AIAgentActionType::CreateDocuments(CreateDocumentsRequest {
documents: required_array(input, "documents")?
.iter()
.enumerate()
.map(|(index, document)| {
require_object(document, &format!("documents[{index}]"))?;
Ok(DocumentToCreate {
title: required_string(document, "title")?,
content: required_string(document, "content")?,
})
})
.collect::<Result<_, String>>()?,
})
}
"edit_plan" | "edit_documents" | "edit_notebook" => {
AIAgentActionType::EditDocuments(EditDocumentsRequest {
diffs: required_array(input, "diffs")?
.iter()
.enumerate()
.map(|(index, diff)| {
require_object(diff, &format!("diffs[{index}]"))?;
let document_id = required_string(diff, "document_id")?;
Ok(DocumentDiff {
document_id: AIDocumentId::try_from(document_id.clone())
.map_err(|_| format!("invalid document_id: {document_id:?}"))?,
search: required_string(diff, "search")?,
replace: required_string(diff, "replace")?,
})
})
.collect::<Result<_, String>>()?,
})
}
"run_agents" => AIAgentActionType::RunAgents(RunAgentsRequest {
summary: required_nonempty_string(input, "summary")?,
base_prompt: optional_string(input, "base_prompt")?.unwrap_or_default(),
skills: skill_references(input, skill_path_origin)?,
model_id: optional_string(input, "model_id")?.unwrap_or_default(),
harness_type: optional_string(input, "harness_type")?.unwrap_or_default(),
execution_mode: run_agents_execution_mode(input)?,
agent_run_configs: nonempty_required_array(input, "agent_run_configs")?
.iter()
.enumerate()
.map(|(index, config)| {
require_object(config, &format!("agent_run_configs[{index}]"))?;
Ok(RunAgentsAgentRunConfig {
name: required_nonempty_string(config, "name")?,
prompt: required_nonempty_string(config, "prompt")?,
title: optional_string(config, "title")?.unwrap_or_default(),
})
})
.collect::<Result<_, String>>()?,
plan_id: optional_string(input, "plan_id")?.unwrap_or_default(),
harness_auth_secret_name: None,
}),
"start_agent" => AIAgentActionType::StartAgent {
version: StartAgentVersion::V1,
name: required_nonempty_string(input, "name")?,
prompt: required_nonempty_string(input, "prompt")?,
execution_mode: StartAgentExecutionMode::local_with_defaults(),
lifecycle_subscription: None,
},
"send_message_to_agent" => AIAgentActionType::SendMessageToAgent {
addresses: vec![required_string(input, "agent_id")?],
subject: String::new(),
message: required_string(input, "message")?,
},
"transfer_shell_command_control_to_user" => {
AIAgentActionType::TransferShellCommandControlToUser {
reason: required_nonempty_string(input, "reason")?,
}
}
"wait_for_events" => AIAgentActionType::WaitForEvents {
tool_call_id: call.id.clone(),
idle_timeout_seconds: optional_nonnegative_i32(input, "idle_timeout_seconds")?
.unwrap_or(0),
},
"ask_user_question" => AIAgentActionType::AskUserQuestion {
questions: vec![AskUserQuestionItem {
question_id: Uuid::new_v4().to_string(),
question: required_string(input, "question")?,
question_type: AskUserQuestionType::MultipleChoice {
is_multiselect: false,
options: optional_strings(input, "options")?
.unwrap_or_default()
.into_iter()
.enumerate()
.map(|(index, label)| AskUserQuestionOption {
label,
recommended: index == 0,
})
.collect(),
supports_other: true,
},
}],
},
"read_skill" => {
let skill = required_string(input, "skill")?;
let skill = match required_string(input, "reference_type")?.as_str() {
"bundled" => SkillReference::BundledSkillId(skill),
"path" => SkillReference::Path(
skill_path_origin
.location_for_path(skill)
.map_err(|error| error.to_string())?,
),
reference_type => {
return Err(format!(
"invalid reference_type: expected \"path\" or \"bundled\", got {reference_type:?}"
));
}
};
AIAgentActionType::ReadSkill(ReadSkillRequest { skill })
}
"fetch_conversation" => AIAgentActionType::FetchConversation {
conversation_id: required_string(input, "conversation_id")?,
},
name if name.starts_with("mcp__") => {
let mut parts = name.splitn(3, "__");
let _prefix = parts.next();
let server_id = parts.next().and_then(|value| Uuid::parse_str(value).ok());
let name = parts
.next()
.unwrap_or_else(|| name.strip_prefix("mcp__").unwrap_or(name))
.to_string();
AIAgentActionType::CallMCPTool {
server_id,
name,
input: input.clone(),
}
}
}
name => return Err(format!("unsupported Rig tool proposal: {name}")),
}
};
@@ -242,155 +268,309 @@ pub(crate) fn action_from_tool_call(
})
}
fn string(input: &serde_json::Value, key: &str) -> String {
fn require_object(input: &serde_json::Value, field: &str) -> Result<(), String> {
input
.get(key)
.and_then(serde_json::Value::as_str)
.unwrap_or_default()
.to_string()
.is_object()
.then_some(())
.ok_or_else(|| format!("invalid {field}: expected an object"))
}
fn nonempty_string(input: &serde_json::Value, key: &str) -> Option<String> {
let value = string(input, key);
(!value.is_empty()).then_some(value)
}
fn boolean(input: &serde_json::Value, key: &str) -> bool {
fn required_string(input: &serde_json::Value, key: &str) -> Result<String, String> {
input
.get(key)
.and_then(serde_json::Value::as_bool)
.unwrap_or(false)
}
fn strings(input: &serde_json::Value, key: &str) -> Vec<String> {
input
.get(key)
.and_then(serde_json::Value::as_array)
.into_iter()
.flatten()
.filter_map(serde_json::Value::as_str)
.ok_or_else(|| format!("missing required field {key:?}"))?
.as_str()
.map(ToOwned::to_owned)
.ok_or_else(|| format!("invalid field {key:?}: expected a string"))
}
fn required_nonempty_string(input: &serde_json::Value, key: &str) -> Result<String, String> {
let value = required_string(input, key)?;
if value.trim().is_empty() {
return Err(format!(
"invalid field {key:?}: expected a non-empty string"
));
}
Ok(value)
}
fn optional_string(input: &serde_json::Value, key: &str) -> Result<Option<String>, String> {
input
.get(key)
.map(|value| {
value
.as_str()
.map(ToOwned::to_owned)
.ok_or_else(|| format!("invalid field {key:?}: expected a string"))
})
.transpose()
}
fn required_array<'a>(
input: &'a serde_json::Value,
key: &str,
) -> Result<&'a Vec<serde_json::Value>, String> {
input
.get(key)
.ok_or_else(|| format!("missing required field {key:?}"))?
.as_array()
.ok_or_else(|| format!("invalid field {key:?}: expected an array"))
}
fn nonempty_required_array<'a>(
input: &'a serde_json::Value,
key: &str,
) -> Result<&'a Vec<serde_json::Value>, String> {
let values = required_array(input, key)?;
if values.is_empty() {
return Err(format!("invalid field {key:?}: expected at least one item"));
}
Ok(values)
}
fn required_strings(input: &serde_json::Value, key: &str) -> Result<Vec<String>, String> {
strings_from_array(required_array(input, key)?, key)
}
fn optional_strings(input: &serde_json::Value, key: &str) -> Result<Option<Vec<String>>, String> {
input
.get(key)
.map(|value| {
let values = value
.as_array()
.ok_or_else(|| format!("invalid field {key:?}: expected an array"))?;
strings_from_array(values, key)
})
.transpose()
}
fn strings_from_array(values: &[serde_json::Value], key: &str) -> Result<Vec<String>, String> {
values
.iter()
.enumerate()
.map(|(index, value)| {
value
.as_str()
.map(ToOwned::to_owned)
.ok_or_else(|| format!("invalid {key}[{index}]: expected a string"))
})
.collect()
}
fn nonempty_strings(input: &serde_json::Value, key: &str) -> Option<Vec<String>> {
let values = strings(input, key);
(!values.is_empty()).then_some(values)
fn required_uuid(input: &serde_json::Value, key: &str) -> Result<Uuid, String> {
let value = required_string(input, key)?;
Uuid::parse_str(&value).map_err(|_| format!("invalid field {key:?}: expected a UUID"))
}
fn uuid(input: &serde_json::Value, key: &str) -> Option<Uuid> {
fn optional_boolean(input: &serde_json::Value, key: &str) -> Result<Option<bool>, String> {
input
.get(key)
.and_then(serde_json::Value::as_str)
.and_then(|value| Uuid::parse_str(value).ok())
.map(|value| {
value
.as_bool()
.ok_or_else(|| format!("invalid field {key:?}: expected a boolean"))
})
.transpose()
}
fn optional_bounded_u64(
input: &serde_json::Value,
key: &str,
maximum: u64,
) -> Result<Option<u64>, String> {
input
.get(key)
.map(|value| {
let value = value
.as_u64()
.ok_or_else(|| format!("invalid field {key:?}: expected a non-negative integer"))?;
if value > maximum {
return Err(format!(
"invalid field {key:?}: expected an integer no greater than {maximum}"
));
}
Ok(value)
})
.transpose()
}
fn optional_nonnegative_i32(input: &serde_json::Value, key: &str) -> Result<Option<i32>, String> {
input
.get(key)
.map(|value| {
value
.as_i64()
.and_then(|value| i32::try_from(value).ok())
.filter(|value| *value >= 0)
.ok_or_else(|| {
format!("invalid field {key:?}: expected a non-negative 32-bit integer")
})
})
.transpose()
}
fn skill_references(
input: &serde_json::Value,
skill_path_origin: &SkillPathOrigin,
) -> Vec<SkillReference> {
input
.get("skills")
.and_then(serde_json::Value::as_array)
.into_iter()
.flatten()
.filter_map(|skill| {
let reference = string(skill, "skill");
if reference.is_empty() {
return None;
}
match skill
.get("reference_type")
.and_then(serde_json::Value::as_str)
{
Some("bundled") => Some(SkillReference::BundledSkillId(reference)),
Some("path") | Some(_) | None => skill_path_origin
) -> Result<Vec<SkillReference>, String> {
let Some(skills) = optional_array(input, "skills")? else {
return Ok(Vec::new());
};
skills
.iter()
.enumerate()
.map(|(index, skill)| {
require_object(skill, &format!("skills[{index}]"))?;
let reference = required_string(skill, "skill")?;
match required_string(skill, "reference_type")?.as_str() {
"bundled" => Ok(SkillReference::BundledSkillId(reference)),
"path" => skill_path_origin
.location_for_path(reference)
.ok()
.map(SkillReference::Path),
.map(SkillReference::Path)
.map_err(|error| error.to_string()),
reference_type => Err(format!(
"invalid skills[{index}].reference_type: expected \"path\" or \"bundled\", got {reference_type:?}"
)),
}
})
.collect()
}
fn run_agents_execution_mode(input: &serde_json::Value) -> RunAgentsExecutionMode {
fn run_agents_execution_mode(input: &serde_json::Value) -> Result<RunAgentsExecutionMode, String> {
let Some(execution_mode) = input.get("execution_mode") else {
return RunAgentsExecutionMode::Local;
return Ok(RunAgentsExecutionMode::Local);
};
let mode_type = execution_mode
.get("type")
.and_then(serde_json::Value::as_str)
.or_else(|| execution_mode.as_str());
match mode_type {
Some("remote") => RunAgentsExecutionMode::Remote {
environment_id: string(execution_mode, "environment_id"),
worker_host: string(execution_mode, "worker_host"),
computer_use_enabled: boolean(execution_mode, "computer_use_enabled"),
},
Some("local") | Some(_) | None => RunAgentsExecutionMode::Local,
require_object(execution_mode, "execution_mode")?;
match optional_string(execution_mode, "type")?.as_deref() {
Some("remote") => Ok(RunAgentsExecutionMode::Remote {
environment_id: optional_string(execution_mode, "environment_id")?.unwrap_or_default(),
worker_host: optional_string(execution_mode, "worker_host")?.unwrap_or_default(),
computer_use_enabled: optional_boolean(execution_mode, "computer_use_enabled")?
.unwrap_or(false),
}),
Some("local") | None => {
optional_string(execution_mode, "environment_id")?;
optional_string(execution_mode, "worker_host")?;
optional_boolean(execution_mode, "computer_use_enabled")?;
Ok(RunAgentsExecutionMode::Local)
}
Some(mode_type) => Err(format!(
"invalid execution_mode.type: expected \"local\" or \"remote\", got {mode_type:?}"
)),
}
}
fn file_location(file: &serde_json::Value) -> Option<FileLocations> {
fn file_location(file: &serde_json::Value, file_index: usize) -> Result<FileLocations, String> {
if let Some(name) = file.as_str() {
return Some(FileLocations {
return Ok(FileLocations {
name: name.to_string(),
lines: Vec::new(),
});
}
let name = file
.get("path")
.or_else(|| file.get("name"))?
.as_str()?
.to_string();
let lines = file
.get("line_ranges")
.and_then(serde_json::Value::as_array)
.into_iter()
.flatten()
.filter_map(|range| {
let start = usize::try_from(range.get("start")?.as_u64()?).ok()?;
let end = usize::try_from(range.get("end")?.as_u64()?).ok()?;
(start > 0 && end >= start).then_some(start..end)
})
.collect();
Some(FileLocations { name, lines })
require_object(file, &format!("files[{file_index}]"))?;
let name = required_string(file, "path")?;
let lines = match file.get("line_ranges") {
None => Vec::new(),
Some(value) => value
.as_array()
.ok_or_else(|| format!("invalid files[{file_index}].line_ranges: expected an array"))?
.iter()
.enumerate()
.map(|(range_index, range)| {
require_object(
range,
&format!("files[{file_index}].line_ranges[{range_index}]"),
)?;
let start = required_line_number(range, "start", file_index, range_index)?;
let inclusive_end = required_line_number(range, "end", file_index, range_index)?;
if inclusive_end < start {
return Err(format!(
"invalid files[{file_index}].line_ranges[{range_index}]: end must be greater than or equal to start"
));
}
let exclusive_end = inclusive_end.checked_add(1).ok_or_else(|| {
format!(
"invalid files[{file_index}].line_ranges[{range_index}].end: inclusive end is too large"
)
})?;
Ok(start..exclusive_end)
})
.collect::<Result<_, String>>()?,
};
Ok(FileLocations { name, lines })
}
fn file_edits(input: &serde_json::Value) -> Vec<FileEdit> {
let diffs = input
.get("diffs")
.and_then(serde_json::Value::as_array)
.into_iter()
.flatten()
.map(|diff| {
FileEdit::Edit(ParsedDiff::StrReplaceEdit {
file: nonempty_string(diff, "file_path"),
search: nonempty_string(diff, "search"),
replace: nonempty_string(diff, "replace"),
})
});
let creates = input
.get("new_files")
.and_then(serde_json::Value::as_array)
.into_iter()
.flatten()
.map(|file| FileEdit::Create {
file: nonempty_string(file, "file_path"),
content: nonempty_string(file, "content"),
});
let deletes = input
.get("deleted_files")
.and_then(serde_json::Value::as_array)
.into_iter()
.flatten()
.map(|file| FileEdit::Delete {
file: file
fn required_line_number(
range: &serde_json::Value,
key: &str,
file_index: usize,
range_index: usize,
) -> Result<usize, String> {
let value = range
.get(key)
.ok_or_else(|| format!("missing required field {key:?}"))?
.as_u64()
.and_then(|value| usize::try_from(value).ok())
.filter(|value| *value > 0)
.ok_or_else(|| {
format!(
"invalid files[{file_index}].line_ranges[{range_index}].{key}: expected a positive integer"
)
})?;
Ok(value)
}
fn file_edits(input: &serde_json::Value) -> Result<Vec<FileEdit>, String> {
let mut edits = Vec::new();
if let Some(diffs) = optional_array(input, "diffs")? {
for (index, diff) in diffs.iter().enumerate() {
require_object(diff, &format!("diffs[{index}]"))?;
edits.push(FileEdit::Edit(ParsedDiff::StrReplaceEdit {
file: Some(required_string(diff, "file_path")?),
search: Some(required_string(diff, "search")?),
replace: Some(required_string(diff, "replace")?),
}));
}
}
if let Some(files) = optional_array(input, "new_files")? {
for (index, file) in files.iter().enumerate() {
require_object(file, &format!("new_files[{index}]"))?;
edits.push(FileEdit::Create {
file: Some(required_string(file, "file_path")?),
content: Some(required_string(file, "content")?),
});
}
}
if let Some(files) = optional_array(input, "deleted_files")? {
for (index, file) in files.iter().enumerate() {
let path = file
.as_str()
.map(ToOwned::to_owned)
.or_else(|| nonempty_string(file, "file_path")),
});
diffs.chain(creates).chain(deletes).collect()
.ok_or_else(|| format!("invalid deleted_files[{index}]: expected a string"))?;
edits.push(FileEdit::Delete {
file: Some(path.to_owned()),
});
}
}
if edits.is_empty() {
return Err(
"invalid file edits: expected at least one diff, new file, or deleted file".to_string(),
);
}
Ok(edits)
}
fn optional_array<'a>(
input: &'a serde_json::Value,
key: &str,
) -> Result<Option<&'a Vec<serde_json::Value>>, String> {
input
.get(key)
.map(|value| {
value
.as_array()
.ok_or_else(|| format!("invalid field {key:?}: expected an array"))
})
.transpose()
}
#[cfg(test)]
app/src/ai/runtime/rig_tool_tests.rs
+241
View File
@@ -46,6 +46,26 @@ fn shell_calls_become_domain_actions_without_a_proto_round_trip() {
));
}
#[test]
fn transfer_control_calls_become_domain_actions() {
let action = action_from_tool_call(
"task-1",
&call(
"transfer_shell_command_control_to_user",
serde_json::json!({"reason": "The command needs interactive input"}),
),
&SkillPathOrigin::Local,
&HashMap::new(),
)
.unwrap();
assert!(matches!(
action.action,
AIAgentActionType::TransferShellCommandControlToUser { reason }
if reason == "The command needs interactive input"
));
}
#[test]
fn create_plan_calls_become_document_actions() {
let action = action_from_tool_call(
@@ -72,6 +92,227 @@ fn create_plan_calls_become_document_actions() {
assert_eq!(request.documents[0].content, "# Implementation plan");
}
#[test]
fn read_files_converts_advertised_inclusive_ranges_to_half_open_ranges() {
let action = action_from_tool_call(
"task-1",
&call(
"read_files",
serde_json::json!({
"files": [{
"path": "/tmp/example.rs",
"line_ranges": [
{"start": 1, "end": 1},
{"start": 10, "end": 25}
]
}]
}),
),
&SkillPathOrigin::Local,
&HashMap::new(),
)
.unwrap();
let AIAgentActionType::ReadFiles(request) = action.action else {
panic!("expected read-files action");
};
assert_eq!(request.locations[0].lines, vec![1..2, 10..26]);
}
#[test]
fn known_tools_reject_malformed_required_inputs() {
let cases = [
("read_files", serde_json::json!({}), "files"),
(
"read_files",
serde_json::json!({"files": "not-an-array"}),
"expected an array",
),
(
"read_files",
serde_json::json!({"files": [{"path": "/tmp/a", "line_ranges": [{"start": 0, "end": 1}]}]}),
"positive integer",
),
(
"read_files",
serde_json::json!({"files": [{"path": "/tmp/a", "line_ranges": [{"start": 3, "end": 2}]}]}),
"greater than or equal",
),
(
"read_files",
serde_json::json!({"files": [{"path": "/tmp/a", "line_ranges": [{"start": 1, "end": u64::MAX}]}]}),
"inclusive end is too large",
),
(
"grep",
serde_json::json!({"queries": ["ok", 7]}),
"queries[1]",
),
(
"file_glob",
serde_json::json!({"patterns": false}),
"expected an array",
),
(
"search_codebase",
serde_json::json!({"query": 42}),
"expected a string",
),
(
"apply_file_diffs",
serde_json::json!({"summary": "edit", "diffs": [{"file_path": "/tmp/a", "search": "x"}]}),
"replace",
),
(
"apply_file_diffs",
serde_json::json!({"summary": "Nothing to do"}),
"at least one diff",
),
(
"run_shell_command",
serde_json::json!({"command": 42}),
"expected a string",
),
(
"run_shell_command",
serde_json::json!({"command": " "}),
"non-empty string",
),
(
"run_shell_command",
serde_json::json!({"command": "pwd", "is_read_only": "yes"}),
"expected a boolean",
),
(
"write_to_long_running_shell_command",
serde_json::json!({"command_id": "command-1", "input": "yes", "mode": "words"}),
"mode",
),
(
"interrupt_shell_command",
serde_json::json!({}),
"command_id",
),
(
"read_shell_command_output",
serde_json::json!({"command_id": 12}),
"expected a string",
),
(
"read_shell_command_output",
serde_json::json!({"command_id": "command-1", "wait_seconds": 11}),
"no greater than",
),
(
"run_agents",
serde_json::json!({"summary": "Investigate", "agent_run_configs": []}),
"at least one item",
),
(
"run_agents",
serde_json::json!({"summary": "Investigate", "agent_run_configs": [{"name": "one"}]}),
"prompt",
),
(
"run_agents",
serde_json::json!({"summary": "Investigate", "agent_run_configs": [{"name": "one", "prompt": "Inspect"}], "execution_mode": {"type": "other"}}),
"execution_mode.type",
),
(
"run_agents",
serde_json::json!({"summary": "Investigate", "agent_run_configs": [{"name": "one", "prompt": "Inspect"}], "skills": [{"skill": "test", "reference_type": "other"}]}),
"skills[0].reference_type",
),
(
"start_agent",
serde_json::json!({"name": "worker"}),
"prompt",
),
(
"transfer_shell_command_control_to_user",
serde_json::json!({"reason": false}),
"expected a string",
),
(
"wait_for_events",
serde_json::json!({"idle_timeout_seconds": -1}),
"non-negative",
),
(
"create_plan",
serde_json::json!({"documents": [{"title": "Plan"}]}),
"content",
),
(
"read_skill",
serde_json::json!({"skill": "/tmp/SKILL.md", "reference_type": "other"}),
"reference_type",
),
(
"fetch_conversation",
serde_json::json!({"conversation_id": null}),
"expected a string",
),
];
for (name, arguments, expected_error) in cases {
let error = action_from_tool_call(
"task-1",
&call(name, arguments),
&SkillPathOrigin::Local,
&HashMap::new(),
)
.unwrap_err();
assert!(
error.contains(expected_error),
"{name} error {error:?} did not contain {expected_error:?}"
);
}
}
#[test]
fn known_tools_preserve_legitimate_optional_defaults() {
let cases = [
("grep", serde_json::json!({"queries": ["needle"]})),
("file_glob", serde_json::json!({"patterns": ["**/*.rs"]})),
(
"ask_user_question",
serde_json::json!({"question": "Continue?"}),
),
(
"apply_file_diffs",
serde_json::json!({"summary": "Create file", "new_files": [{"file_path": "/tmp/new", "content": ""}]}),
),
("run_shell_command", serde_json::json!({"command": "pwd"})),
(
"write_to_long_running_shell_command",
serde_json::json!({"command_id": "command-1", "input": ""}),
),
(
"read_shell_command_output",
serde_json::json!({"command_id": "command-1"}),
),
(
"run_agents",
serde_json::json!({
"summary": "Investigate",
"agent_run_configs": [{"name": "worker", "prompt": "Inspect"}]
}),
),
("wait_for_events", serde_json::json!({})),
];
for (name, arguments) in cases {
action_from_tool_call(
"task-1",
&call(name, arguments),
&SkillPathOrigin::Local,
&HashMap::new(),
)
.unwrap_or_else(|error| panic!("{name} rejected optional defaults: {error}"));
}
}
#[test]
fn edit_calls_preserve_file_edits_in_the_domain_model() {
let action = action_from_tool_call(