fix: harden ticket role launch execution

This commit is contained in:
Keisuke Hirata 2026-06-06 04:29:47 +09:00
parent 4bf0e2715c
commit dd70517f96
No known key found for this signature in database
2 changed files with 70 additions and 11 deletions

View File

@ -22,7 +22,7 @@ use uuid::Uuid;
const READY_PREFIX: &str = "YOI-READY\t"; const READY_PREFIX: &str = "YOI-READY\t";
const READY_TIMEOUT: Duration = Duration::from_secs(20); const READY_TIMEOUT: Duration = Duration::from_secs(20);
/// `spawn_pod` の入力。 #[derive(Debug, Clone, PartialEq, Eq)]
pub struct SpawnConfig { pub struct SpawnConfig {
pub runtime_command: PodRuntimeCommand, pub runtime_command: PodRuntimeCommand,
/// `pod.name` として使う識別子。runtime ディレクトリ /// `pod.name` として使う識別子。runtime ディレクトリ

View File

@ -6,8 +6,9 @@
use std::io; use std::io;
use std::path::PathBuf; use std::path::PathBuf;
use std::time::Duration;
use protocol::{Method, Segment}; use protocol::{ErrorCode, Event, InvokeKind, Method, Segment};
use thiserror::Error; use thiserror::Error;
use ticket::config::{TicketConfig, TicketConfigError, TicketRole}; use ticket::config::{TicketConfig, TicketConfigError, TicketRole};
@ -15,6 +16,7 @@ use crate::{PodClient, PodRuntimeCommand, SpawnConfig, SpawnError, SpawnReady, s
const MAX_FIELD_CHARS: usize = 8_000; const MAX_FIELD_CHARS: usize = 8_000;
const MAX_POD_NAME_CHARS: usize = 80; const MAX_POD_NAME_CHARS: usize = 80;
const RUN_ACCEPTANCE_TIMEOUT: Duration = Duration::from_secs(10);
/// Ticket identifier carried by a role launch request. /// Ticket identifier carried by a role launch request.
#[derive(Debug, Clone, PartialEq, Eq, Default)] #[derive(Debug, Clone, PartialEq, Eq, Default)]
@ -119,15 +121,21 @@ impl TicketRoleLaunchPlan {
} }
} }
pub fn spawn_config(&self, runtime_command: PodRuntimeCommand) -> SpawnConfig { pub fn spawn_config(
SpawnConfig { &self,
runtime_command: PodRuntimeCommand,
) -> Result<SpawnConfig, TicketRoleLaunchError> {
if self.profile == "inherit" {
return Err(TicketRoleLaunchError::UnsupportedInheritProfile);
}
Ok(SpawnConfig {
runtime_command, runtime_command,
pod_name: self.pod_name.clone(), pod_name: self.pod_name.clone(),
profile: Some(self.profile.clone()), profile: Some(self.profile.clone()),
cwd: self.workspace_root.clone(), cwd: self.workspace_root.clone(),
resume_from: None, resume_from: None,
resume_by_pod_name: false, resume_by_pod_name: false,
} })
} }
} }
@ -144,6 +152,10 @@ pub enum TicketRoleLaunchError {
Config(#[from] TicketConfigError), Config(#[from] TicketConfigError),
#[error("Ticket role Pod name must not be empty")] #[error("Ticket role Pod name must not be empty")]
EmptyPodName, EmptyPodName,
#[error(
"Ticket role profile 'inherit' cannot be used for top-level launch execution; configure a concrete role profile selector"
)]
UnsupportedInheritProfile,
#[error(transparent)] #[error(transparent)]
Spawn(#[from] SpawnError), Spawn(#[from] SpawnError),
#[error("failed to connect to spawned Ticket role Pod at {}: {source}", .socket_path.display())] #[error("failed to connect to spawned Ticket role Pod at {}: {source}", .socket_path.display())]
@ -157,6 +169,12 @@ pub enum TicketRoleLaunchError {
#[source] #[source]
source: io::Error, source: io::Error,
}, },
#[error("Ticket role Pod rejected first run input with {code:?}: {message}")]
RunRejected { code: ErrorCode, message: String },
#[error("Ticket role Pod closed before confirming first run acceptance")]
RunAcceptanceClosed,
#[error("timed out waiting for Ticket role Pod to confirm first run acceptance")]
RunAcceptanceTimeout,
} }
/// Load `.yoi/ticket.config.toml` from the workspace and construct a launch plan. /// Load `.yoi/ticket.config.toml` from the workspace and construct a launch plan.
@ -202,7 +220,8 @@ pub fn plan_ticket_role_launch_with_config(
}) })
} }
/// Spawn the Pod, connect to its socket, and send the first `Method::Run` input. /// Spawn the Pod, connect to its socket, send the first `Method::Run` input,
/// and wait for bounded acceptance evidence from the Pod event stream.
pub async fn launch_ticket_role_pod<F>( pub async fn launch_ticket_role_pod<F>(
context: TicketRoleLaunchContext, context: TicketRoleLaunchContext,
runtime_command: PodRuntimeCommand, runtime_command: PodRuntimeCommand,
@ -212,7 +231,7 @@ where
F: FnMut(&str), F: FnMut(&str),
{ {
let plan = plan_ticket_role_launch(context)?; let plan = plan_ticket_role_launch(context)?;
let ready = spawn_pod(plan.spawn_config(runtime_command), progress).await?; let ready = spawn_pod(plan.spawn_config(runtime_command)?, progress).await?;
let mut client = PodClient::connect(&ready.socket_path) let mut client = PodClient::connect(&ready.socket_path)
.await .await
.map_err(|source| TicketRoleLaunchError::Connect { .map_err(|source| TicketRoleLaunchError::Connect {
@ -223,9 +242,39 @@ where
.send(&plan.run_method()) .send(&plan.run_method())
.await .await
.map_err(|source| TicketRoleLaunchError::SendRun { source })?; .map_err(|source| TicketRoleLaunchError::SendRun { source })?;
wait_for_run_acceptance(&mut client, &plan.run_segments, RUN_ACCEPTANCE_TIMEOUT).await?;
Ok(TicketRoleLaunchResult { plan, ready }) Ok(TicketRoleLaunchResult { plan, ready })
} }
async fn wait_for_run_acceptance(
client: &mut PodClient,
expected_segments: &[Segment],
timeout: Duration,
) -> Result<(), TicketRoleLaunchError> {
let wait = async {
loop {
let Some(event) = client.next_event().await else {
return Err(TicketRoleLaunchError::RunAcceptanceClosed);
};
match event {
Event::UserMessage { segments } if segments == expected_segments => return Ok(()),
Event::InvokeStart {
kind: InvokeKind::UserSend,
}
| Event::TurnStart { .. } => return Ok(()),
Event::Error { code, message } => {
return Err(TicketRoleLaunchError::RunRejected { code, message });
}
_ => {}
}
}
};
tokio::time::timeout(timeout, wait)
.await
.map_err(|_| TicketRoleLaunchError::RunAcceptanceTimeout)?
}
fn build_launch_prompt( fn build_launch_prompt(
context: &TicketRoleLaunchContext, context: &TicketRoleLaunchContext,
profile: &str, profile: &str,
@ -414,10 +463,14 @@ mod tests {
Segment::WorkflowInvoke { slug } if slug == "multi-agent-workflow" Segment::WorkflowInvoke { slug } if slug == "multi-agent-workflow"
)); ));
assert!(text_segment(&plan).contains("Profile selector: inherit")); assert!(text_segment(&plan).contains("Profile selector: inherit"));
let spawn = plan.spawn_config(PodRuntimeCommand::for_executable("/bin/yoi")); let err = plan
assert_eq!(spawn.pod_name, "ticket-coder-ticket-role-pod-launcher"); .spawn_config(PodRuntimeCommand::for_executable("/bin/yoi"))
assert_eq!(spawn.profile.as_deref(), Some("inherit")); .unwrap_err();
assert_eq!(spawn.cwd, temp.path()); assert!(matches!(
err,
TicketRoleLaunchError::UnsupportedInheritProfile
));
assert!(err.to_string().contains("'inherit' cannot be used"));
} }
#[test] #[test]
@ -460,6 +513,12 @@ workflow = "ticket-review-workflow"
assert!(text.contains("Workflow: ticket-review-workflow")); assert!(text.contains("Workflow: ticket-review-workflow"));
assert!(text.contains("Profile selector: project:reviewer")); assert!(text.contains("Profile selector: project:reviewer"));
assert!(!text.contains("system_instruction")); assert!(!text.contains("system_instruction"));
let spawn = plan
.spawn_config(PodRuntimeCommand::for_executable("/bin/yoi"))
.unwrap();
assert_eq!(spawn.pod_name, "reviewer-fixed");
assert_eq!(spawn.profile.as_deref(), Some("project:reviewer"));
assert_eq!(spawn.cwd, temp.path());
} }
#[test] #[test]