fix(server): persist cancelled terminal state and align status tests

Persist a cancelled terminal record when a live run is interrupted by the
server-side cancel signal, and abort pending web interview questions so
human-gated runs can unwind instead of hanging in a non-terminal durable
state.

Also align server tests with the current succeeded status contract and poll
aggregate usage until the in-memory accumulator converges with the store-
backed run status.
This commit is contained in:
Bryan Helmkamp 2026-04-05 10:42:22 -04:00
parent 21671ffa0a
commit 0a9e583ce5
No known key found for this signature in database
3 changed files with 77 additions and 24 deletions

View file

@ -77,6 +77,19 @@ impl WebInterviewer {
};
sender.is_some_and(|tx| tx.send(answer).is_ok())
}
/// Abort all pending questions, unblocking any waiting `ask()` calls.
pub fn abort_pending(&self) {
let pending = {
let mut inner = self.inner.lock().expect("web interviewer lock poisoned");
inner.questions.clear();
inner.pending.drain().map(|(_, tx)| tx).collect::<Vec<_>>()
};
for sender in pending {
let _ = sender.send(Answer::aborted());
}
}
}
impl Default for WebInterviewer {
@ -260,6 +273,24 @@ mod tests {
assert!(interviewer.pending_questions().is_empty());
}
#[tokio::test]
async fn abort_pending_unblocks_ask_and_clears_questions() {
let interviewer = Arc::new(WebInterviewer::new());
let interviewer_clone = Arc::clone(&interviewer);
let ask_handle = tokio::spawn(async move {
let q = Question::new("approve?", QuestionType::YesNo);
interviewer_clone.ask(q).await
});
wait_for_pending_count(interviewer.as_ref(), 1).await;
interviewer.abort_pending();
let answer = ask_handle.await.expect("task should complete");
assert_eq!(answer.value, AnswerValue::Aborted);
assert!(interviewer.pending_questions().is_empty());
}
#[tokio::test]
async fn ask_returns_aborted_when_pending_sender_is_dropped() {
let interviewer = Arc::new(WebInterviewer::new());

View file

@ -146,6 +146,11 @@ enum RunExecutionMode {
Resume,
}
enum ExecutionResult {
Completed(Result<operations::Started, FabroError>),
CancelledBySignal,
}
/// Per-model usage totals.
#[derive(Default)]
struct ModelUsageTotals {
@ -1080,13 +1085,19 @@ async fn execute_run(state: Arc<AppState>, run_id: RunId) {
};
let result = tokio::select! {
result = execution => result,
result = execution => ExecutionResult::Completed(result),
_ = cancel_rx => {
cancel_token.store(true, Ordering::SeqCst);
Err(FabroError::Cancelled)
ExecutionResult::CancelledBySignal
}
};
if matches!(result, ExecutionResult::CancelledBySignal) {
if let Err(err) = persist_cancelled_run_status(state.as_ref(), run_id).await {
error!(run_id = %run_id, error = %err, "Failed to persist cancelled run status");
}
}
// Save final checkpoint
let checkpoint = match run_store.state().await {
Ok(state) => state.checkpoint,
@ -1128,7 +1139,7 @@ async fn execute_run(state: Arc<AppState>, run_id: RunId) {
let mut runs = state.runs.lock().expect("runs lock poisoned");
if let Some(managed_run) = runs.get_mut(&run_id) {
match &result {
Ok(started) => match &started.finalized.outcome {
ExecutionResult::Completed(Ok(started)) => match &started.finalized.outcome {
Ok(_) => {
info!(run_id = %run_id, "Run completed");
managed_run.status = RunStatus::Completed;
@ -1143,11 +1154,12 @@ async fn execute_run(state: Arc<AppState>, run_id: RunId) {
managed_run.error = Some(e.to_string());
}
},
Err(FabroError::Cancelled) => {
ExecutionResult::Completed(Err(FabroError::Cancelled))
| ExecutionResult::CancelledBySignal => {
info!(run_id = %run_id, "Run cancelled");
managed_run.status = RunStatus::Cancelled;
}
Err(e) => {
ExecutionResult::Completed(Err(e)) => {
error!(run_id = %run_id, error = %e, "Run failed");
managed_run.status = RunStatus::Failed;
managed_run.error = Some(e.to_string());
@ -1695,6 +1707,9 @@ async fn cancel_run(
if let Some(token) = &managed_run.cancel_token {
token.store(true, Ordering::Relaxed);
}
if let Some(interviewer) = &managed_run.interviewer {
interviewer.abort_pending();
}
if let Some(cancel_tx) = managed_run.cancel_tx.take() {
let _ = cancel_tx.send(());
}
@ -3027,11 +3042,11 @@ mod tests {
assert_eq!(response.status(), StatusCode::OK);
let body = body_json(response.into_body()).await;
status = body["status"].as_str().unwrap().to_string();
if status == "completed" || status == "failed" {
if status == "succeeded" || status == "failed" {
break;
}
}
assert_eq!(status, "completed");
assert_eq!(status, "succeeded");
}
#[tokio::test]
@ -3226,24 +3241,31 @@ mod tests {
let response = app.clone().oneshot(req).await.unwrap();
let body = body_json(response.into_body()).await;
status = body["status"].as_str().unwrap().to_string();
if status == "completed" || status == "failed" {
if status == "succeeded" || status == "failed" {
break;
}
}
assert_eq!(status, "completed");
assert_eq!(status, "succeeded");
// Check aggregate usage
let req = Request::builder()
.method("GET")
.uri(api("/usage"))
.body(Body::empty())
.unwrap();
let mut total_runs = 0;
for _ in 0..POLL_ATTEMPTS {
let req = Request::builder()
.method("GET")
.uri(api("/usage"))
.body(Body::empty())
.unwrap();
let response = app.oneshot(req).await.unwrap();
assert_eq!(response.status(), StatusCode::OK);
let response = app.clone().oneshot(req).await.unwrap();
assert_eq!(response.status(), StatusCode::OK);
let body = body_json(response.into_body()).await;
assert_eq!(body["totals"]["runs"].as_i64().unwrap(), 1);
let body = body_json(response.into_body()).await;
total_runs = body["totals"]["runs"].as_i64().unwrap();
if total_runs == 1 {
break;
}
tokio::time::sleep(POLL_INTERVAL).await;
}
assert_eq!(total_runs, 1);
}
#[tokio::test]

View file

@ -572,9 +572,9 @@ mod server_lifecycle {
let response = app.clone().oneshot(req).await.unwrap();
assert_eq!(response.status(), StatusCode::NO_CONTENT);
// 4. Poll until completed
let final_status = wait_for_run_status(&app, &run_id, &["completed", "failed"]).await;
assert_eq!(final_status, "completed");
// 4. Poll until the run reaches a terminal success or failure state.
let final_status = wait_for_run_status(&app, &run_id, &["succeeded", "failed"]).await;
assert_eq!(final_status, "succeeded");
// 5. Verify no pending questions
let req = Request::builder()
@ -917,8 +917,8 @@ mod serve_dry_run {
let response = app.clone().oneshot(req).await.unwrap();
assert_eq!(response.status(), StatusCode::OK);
let status = wait_for_run_status(&app, &run_id, &["completed", "failed"]).await;
assert_eq!(status, "completed");
let status = wait_for_run_status(&app, &run_id, &["succeeded", "failed"]).await;
assert_eq!(status, "succeeded");
}
#[tokio::test]