Skip to content

fix(frontend): keep completion errors across full aggregation #12786

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
KrishnanPrash wants to merge 1 commit into from
Closed

fix(frontend): keep completion errors across full aggregation #12786

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
187 changes: 40 additions & 147 deletions lib/llm/src/http/service/openai.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -53,6 +53,7 @@ use crate::protocols::common::extensions::{
};
use crate::protocols::openai::chat_completions::aggregator::ChatCompletionAggregator;
use crate::protocols::openai::{
ParsingOptions,
audios::{NvAudioSpeechResponse, NvCreateAudioSpeechRequest},
chat_completions::{
NvCreateChatCompletionRequest, NvCreateChatCompletionResponse,
Expand Down Expand Up @@ -849,17 +850,6 @@ async fn completions_single(

Ok(sse_stream.into_response())
} else {
// Preserve typed backend errors before the completions aggregator turns
// them into strings. In particular, Python ValueError/TypeError arrives
// as Backend(InvalidArgument) and must remain an HTTP 400.
let stream = check_for_backend_error(stream, None)
.await
.map_err(|error_response| {
tracing::error!(request_id, "Backend error detected: {:?}", error_response);
inflight_guard.mark_error(extract_error_type_from_response(&error_response));
error_response
})?;

// Tap the stream to collect metrics for non-streaming requests without altering items
let mut http_queue_guard = Some(http_queue_guard);
let stream = stream.inspect(move |response| {
Expand All @@ -871,19 +861,10 @@ async fn completions_single(
);
});

let response = NvCreateCompletionResponse::from_annotated_stream(stream, parsing_options)
let response = aggregate_completion_response(stream, parsing_options, &request_id)
.await
.map_err(|e| {
tracing::error!(
"Failed to fold completions stream for {}: {:?}",
request_id,
e
);
let err_response = ErrorMessage::internal_server_error(&format!(
"Failed to fold completions stream for {request_id}"
));
inflight_guard.mark_error(extract_error_type_from_response(&err_response));
err_response
.inspect_err(|error_response| {
inflight_guard.mark_error(extract_error_type_from_response(error_response));
})?;

inflight_guard.mark_ok();
Expand Down Expand Up @@ -976,26 +957,40 @@ fn aggregate_batch_completion_usage(
}
}

type BoxedCompletionResponseStream =
std::pin::Pin<Box<dyn futures::Stream<Item = Annotated<NvCreateCompletionResponse>> + Send>>;

/// Check each prompt stream before merging a non-streaming completion batch.
/// Fold a non-streaming completion stream into a single response.
///
/// `select_all` cannot safely provide this check after merging because a normal
/// event from one prompt may arrive before a typed backend error from another.
/// Poll all streams concurrently so batch startup is not serialized.
async fn check_completion_batch_streams<S>(
streams: Vec<S>,
) -> Result<Vec<BoxedCompletionResponseStream>, ErrorResponse>
where
S: futures::Stream<Item = Annotated<NvCreateCompletionResponse>> + Send + 'static,
{
futures::future::try_join_all(
streams
.into_iter()
.map(|stream| check_for_backend_error(stream, None)),
)
.await
/// The aggregator reduces a backend error to a string, so the typed error is
/// captured from the raw events and returned with its own status. Every event
/// is inspected, so a backend that fails after its first chunk is covered too.
async fn aggregate_completion_response(
stream: impl futures::Stream<Item = Annotated<NvCreateCompletionResponse>>,
parsing_options: ParsingOptions,
request_id: &str,
) -> Result<NvCreateCompletionResponse, ErrorResponse> {
let backend_error = Arc::new(std::sync::OnceLock::new());
let first_error = backend_error.clone();
let stream = stream.inspect(move |response| {
if let Some(error) = extract_backend_error_if_present(response) {

@dynamo-review-agent dynamo-review-agent Bot Aug 6, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Every non-streaming completion chunk now passes through extract_backend_error_if_present, whose normal-data path serializes NvCreateCompletionResponse to JSON even though that type cannot be a top-level {message, code} error payload, adding avoidable per-token CPU and allocations. Fix: use a completion-specific error extractor that skips generic data-payload JSON conversion for normal completion frames.

🤖 AI Fix

In lib/llm/src/http/service/openai.rs, add extract_completion_backend_error_if_present(event: &Annotated<NvCreateCompletionResponse>) that preserves typed event == "error" handling but does not run serde_json::to_value on event.data, and call it from aggregate_completion_response.

// Keep the first error; it is the one that ended generation.
let _ = first_error.set(error);
}
});

let aggregated =
NvCreateCompletionResponse::from_annotated_stream(stream, parsing_options).await;

@dynamo-review-agent dynamo-review-agent Bot Aug 6, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

aggregate_completion_response records backend errors but waits for NvCreateCompletionResponse::from_annotated_stream to drain the entire merged stream before returning them, so a batch request with one immediate backend error keeps the other prompt streams generating and delays the 4xx response. Fix: short-circuit on the first backend error and drop the remaining stream instead of checking backend_error only after aggregation completes.

🤖 AI Fix

In lib/llm/src/http/service/openai.rs, rewrite aggregate_completion_response to poll the completion stream directly, return Err(backend_error_response(...)) immediately when a backend error frame is seen, and otherwise aggregate non-error Annotated<NvCreateCompletionResponse> items to the same final response.


if let Some((message, status)) = backend_error.get() {
let error_response = backend_error_response(message.clone(), *status);
tracing::warn!(request_id, ?error_response, "Backend error detected");
return Err(error_response);
}

aggregated.map_err(|e| {
tracing::error!(request_id, "Failed to fold completions stream: {e:?}");
ErrorMessage::internal_server_error(&format!(
"Failed to fold completions stream for {request_id}"
))
})
Comment on lines +965 to +993

@devin-ai-integration devin-ai-integration Bot Aug 6, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🔍 Error frames no longer short-circuit; a backend that errors but keeps the stream open will hang the request

The removed preflight (check_for_backend_error) returned as soon as the first non-annotation event was an error frame, so the HTTP response was produced without waiting for the stream to terminate. aggregate_completion_response (lib/llm/src/http/service/openai.rs:965-993) only inspects the captured error after from_annotated_stream has folded the entire stream to completion. If a backend emits a typed error frame and then fails to close the stream (or is slow to do so), the client now blocks until the stream ends instead of getting an immediate 4xx. This is fine for backends where an error frame is terminal (which is what the new tests model), but it is worth confirming that all worker paths close the stream right after emitting event: "error".

Open in Devin Review

Was this helpful? React with 👍 or 👎 to provide feedback.

}

/// Handle batch prompt completions (multiple prompts with n choices each)
Expand Down Expand Up @@ -1100,23 +1095,6 @@ async fn completions_batch(
all_streams.push(remapped_stream);
}

let all_streams: Vec<BoxedCompletionResponseStream> = if streaming {
all_streams
.into_iter()
.map(|stream| Box::pin(stream) as BoxedCompletionResponseStream)
.collect()
} else {
check_completion_batch_streams(all_streams)
.await
.map_err(|error_response| {
tracing::error!(request_id, "Backend error detected: {:?}", error_response);
inflight_guard.mark_error(extract_error_type_from_response(&error_response));
error_response
})?
};

// Merge all streams after every non-streaming prompt has passed its own
// backend-error preflight.
let merged_stream = stream::select_all(all_streams);
let merged_stream = aggregate_batch_completion_usage(merged_stream, request_id.clone());

Expand Down Expand Up @@ -1189,19 +1167,10 @@ async fn completions_batch(
);
});

let response = NvCreateCompletionResponse::from_annotated_stream(stream, parsing_options)
let response = aggregate_completion_response(stream, parsing_options, &request_id)
.await
.map_err(|e| {
tracing::error!(
"Failed to fold completions stream for {}: {:?}",
request_id,
e
);
let err_response = ErrorMessage::internal_server_error(&format!(
"Failed to fold completions stream for {request_id}"
));
inflight_guard.mark_error(extract_error_type_from_response(&err_response));
err_response
.inspect_err(|error_response| {
inflight_guard.mark_error(extract_error_type_from_response(error_response));
})?;

inflight_guard.mark_ok();
Expand Down Expand Up @@ -5128,82 +5097,6 @@ mod tests {
}
}

#[tokio::test]
async fn test_completion_backend_invalid_argument_surfaces_as_400() {
use dynamo_runtime::error::{BackendError, DynamoError, ErrorType};
use futures::stream;

let error_event = Annotated::<NvCreateCompletionResponse> {
data: None,
id: None,
event: Some("error".to_string()),
comment: None,
error: Some(
DynamoError::builder()
.error_type(ErrorType::Backend(BackendError::InvalidArgument))
.message("Dynamo's SGLang backend does not currently support logprobs >= 1")
.build(),
),
};

let error_response =
match check_for_backend_error(stream::iter(vec![error_event]), None).await {
Ok(_) => panic!("typed completion error must fail"),
Err(error_response) => error_response,
};

assert_eq!(error_response.0, StatusCode::BAD_REQUEST);
assert_eq!(error_response.1.code, StatusCode::BAD_REQUEST.as_u16());
assert_eq!(error_response.1.error_type, "Bad Request");
assert!(
error_response
.1
.message
.contains("does not currently support logprobs >= 1")
);
}

#[tokio::test]
async fn test_batch_completion_checks_every_stream_for_backend_errors() {
use dynamo_runtime::error::{BackendError, DynamoError, ErrorType};
use futures::stream;

let normal_event = Annotated::<NvCreateCompletionResponse> {
data: Some(make_completion_chunk("ok", None, None)),
id: None,
event: None,
comment: None,
error: None,
};
let error_event = Annotated::<NvCreateCompletionResponse> {
data: None,
id: None,
event: Some("error".to_string()),
comment: None,
error: Some(
DynamoError::builder()
.error_type(ErrorType::Backend(BackendError::InvalidArgument))
.message("invalid second prompt")
.build(),
),
};

let result = check_completion_batch_streams(vec![
stream::iter(vec![normal_event]),
stream::iter(vec![error_event]),
])
.await;

let error_response = match result {
Ok(_) => panic!("an error in any batch prompt must fail the request"),
Err(error_response) => error_response,
};
assert_eq!(error_response.0, StatusCode::BAD_REQUEST);
assert_eq!(error_response.1.code, StatusCode::BAD_REQUEST.as_u16());
assert_eq!(error_response.1.error_type, "Bad Request");
assert_eq!(error_response.1.message, "invalid second prompt");
}

#[tokio::test]
async fn test_check_for_backend_error_with_json_error_and_code() {
use crate::types::openai::chat_completions::NvCreateChatCompletionStreamResponse;
Expand Down
Loading
Loading