Handle incomplete Gemini streams as provider failures (#8706)

This commit is contained in:
2026-07-18 12:28:43 -07:00
committed by GitHub
parent 0c0ff85513
commit da0d1c04fd
9 changed files with 213 additions and 70 deletions
@@ -64,6 +64,7 @@ scala_library(
],
deps = [
":api_keys",
":external_text_generation_error",
":external_text_generation_service_impl",
":rate_limits",
":streaming_text_results",
@@ -186,16 +186,16 @@ final class ExternalTextGenerationCaller(
case Success(_) =>
if !completed.get() then {
val duration = System.currentTimeMillis() - startTime
println(
val message =
s"[LLM STREAM DIAGNOSTIC] SSE future succeeded without completed callback: " +
s"providerImpl=${serviceImpl.getClass.getSimpleName}, durationMillis=$duration, " +
s"inputChars=${inputText.length}, partialChars=$partialCompletionChars, " +
s"parsedCallbacks=${parsedCallbacks.get()}, terminalCallbacks=${terminalCallbacks.get()}, " +
s"parsedChars=${parsedChars.get()}, lastStreamId=${lastStreamId.get()}, " +
s"lastCallbackCompleted=${lastCallbackCompleted.get()}, sse={${sseListener.diagnostics}}"
)
}
Success(())
println(message)
Failure(ExternalTextGenerationError.ServerError(code = 502, msg = message))
} else Success(())
case Failure(e: SocketTimeoutException) =>
val duration = System.currentTimeMillis() - startTime
@@ -213,6 +213,9 @@ final class ExternalTextGenerationCaller(
)
)
case Failure(e: ExternalTextGenerationStreamFailure) =>
Failure(ExternalTextGenerationError.ServerError(code = 502, msg = e.getMessage))
case Failure(exception) => Failure(exception)
}.recoverWith {
// Don't retry 4xx client errors - they won't succeed
@@ -1,5 +1,8 @@
package net.eagle0.common.llm_integration
private[llm_integration] final class ExternalTextGenerationStreamFailure(message: String)
extends RuntimeException(message)
enum ExternalTextGenerationError extends Error:
case RateLimit(code: Int, msg: String)
case Http(code: Int, msg: String)
@@ -107,6 +107,12 @@ class GeminiServiceImpl(
private var promptTokenCount = 0
private var candidatesTokenCount = 0
private def failStream(message: String): Nothing =
throw new ExternalTextGenerationStreamFailure(message)
private def optionalString(parsedJson: org.json4s.JValue, fieldName: String): Option[String] =
(parsedJson \ fieldName).extractOpt[String].filter(_.nonEmpty)
override def accept(t: String): Unit = {
val parsedJsonOpt =
try
@@ -118,86 +124,98 @@ class GeminiServiceImpl(
}
parsedJsonOpt.foreach { parsedJson =>
val hasError = (parsedJson \ "error") match {
val apiErrorMessage = (parsedJson \ "error") match {
case JObject(errorFields) if errorFields.nonEmpty =>
val errorMessage = (parsedJson \ "error" \ "message").extractOpt[String].getOrElse("Unknown error")
println(s"Gemini API error: $errorMessage")
true
case _ => false
Some((parsedJson \ "error" \ "message").extractOpt[String].getOrElse("Unknown error"))
case _ => None
}
if !hasError then {
// Extract usageMetadata if present (Gemini includes this in streaming responses)
(parsedJson \ "usageMetadata" \ "promptTokenCount").extractOpt[Int].foreach { count =>
promptTokenCount = count
}
(parsedJson \ "usageMetadata" \ "candidatesTokenCount").extractOpt[Int].foreach { count =>
candidatesTokenCount = count
}
apiErrorMessage.foreach(message => failStream(s"Gemini API error: $message"))
val candidatesOpt = (parsedJson \ "candidates") match {
case JArray(arr) => Some(arr)
case _ => None
}
val promptBlockReason = optionalString(parsedJson \ "promptFeedback", "blockReason")
.filterNot(_ == "BLOCK_REASON_UNSPECIFIED")
promptBlockReason.foreach { reason =>
val blockMessage = optionalString(parsedJson \ "promptFeedback", "blockReasonMessage")
failStream(
s"Gemini blocked the prompt with blockReason=$reason" +
blockMessage.fold("")(message => s", message=$message")
)
}
candidatesOpt.filter(_.nonEmpty).foreach { candidates =>
val firstCandidate =
candidates.headOption match {
case Some(obj: JObject) => obj
case Some(other) =>
throw new IllegalArgumentException(s"Expected Gemini candidate object, got $other")
case None =>
throw new IllegalArgumentException("Expected Gemini candidate object, got no candidates")
}
// Extract usageMetadata if present (Gemini includes this in streaming responses)
(parsedJson \ "usageMetadata" \ "promptTokenCount").extractOpt[Int].foreach { count =>
promptTokenCount = count
}
(parsedJson \ "usageMetadata" \ "candidatesTokenCount").extractOpt[Int].foreach { count =>
candidatesTokenCount = count
}
// Check finish reason
val finishReason = (firstCandidate \ "finishReason") match {
case JString(reason) => Some(reason)
case _ => None
val candidatesOpt = (parsedJson \ "candidates") match {
case JArray(arr) => Some(arr)
case _ => None
}
candidatesOpt.filter(_.nonEmpty).foreach { candidates =>
val firstCandidate =
candidates.headOption match {
case Some(obj: JObject) => obj
case Some(other) =>
throw new IllegalArgumentException(s"Expected Gemini candidate object, got $other")
case None =>
throw new IllegalArgumentException("Expected Gemini candidate object, got no candidates")
}
val completed = finishReason.exists(r => r == "STOP" || r == "MAX_TOKENS" || r == "SAFETY")
// Check finish reason
val finishReason = optionalString(firstCandidate, "finishReason")
val finishMessage = optionalString(firstCandidate, "finishMessage")
val completed = finishReason.exists(reason => reason == "STOP" || reason == "MAX_TOKENS")
// Extract text content
val textContent = for {
content <- (firstCandidate \ "content").toOption
parts <- (content \ "parts") match {
case JArray(arr) => Some(arr)
case _ => None
}
firstPart <- parts.headOption.collect { case obj: JObject => obj }
text <- (firstPart \ "text") match {
case JString(s) => Some(s)
case _ => None
}
} yield text
finishReason.filterNot(reason => reason == "STOP" || reason == "MAX_TOKENS").foreach { reason =>
failStream(
s"Gemini stopped generation with finishReason=$reason" +
finishMessage.fold("")(message => s", message=$message")
)
}
val tokenUsage =
if completed then Some(TokenUsage(inputTokens = promptTokenCount, outputTokens = candidatesTokenCount))
else None
// Extract text content
val textContent = for {
content <- (firstCandidate \ "content").toOption
parts <- (content \ "parts") match {
case JArray(arr) => Some(arr)
case _ => None
}
firstPart <- parts.headOption.collect { case obj: JObject => obj }
text <- (firstPart \ "text") match {
case JString(s) => Some(s)
case _ => None
}
} yield text
textContent match {
case Some(text) =>
val tokenUsage =
if completed then Some(TokenUsage(inputTokens = promptTokenCount, outputTokens = candidatesTokenCount))
else None
textContent match {
case Some(text) =>
streamingConsumer.accept(
StreamingTextResults(
streamId = streamId,
value = text,
completed = completed,
tokenUsage = tokenUsage
)
)
case None =>
// No text content but might be completed
if completed then
streamingConsumer.accept(
StreamingTextResults(
streamId = streamId,
value = text,
completed = completed,
value = "",
completed = true,
tokenUsage = tokenUsage
)
)
case None =>
// No text content but might be completed
if completed then
streamingConsumer.accept(
StreamingTextResults(
streamId = streamId,
value = "",
completed = true,
tokenUsage = tokenUsage
)
)
}
}
}
}
@@ -141,6 +141,15 @@ object LlmResolver {
): Boolean =
streamedChars == 0 && partialCompletion.forall(_.isEmpty)
private[service] def incompleteStreamError(
providerName: String,
requestSummary: String
): ExternalTextGenerationError.ServerError =
ExternalTextGenerationError.ServerError(
code = 502,
msg = s"Provider $providerName closed its stream without a terminal callback for $requestSummary"
)
case class LlmRequestWithGameState(
llmRequest: GeneratedTextRequestT,
gameState: GameState,
@@ -613,10 +622,12 @@ class LlmResolver(updateReceiver: LlmUpdateQueuingProxy, isGameValid: GameId =>
Success(result)
} else {
logProviderFutureFinished("succeeded-without-completed-callback")
recordProviderFailure(providerName)
LlmUsageTracker.recordFailure(providerName, llmRequest.getClass.getSimpleName)
Failure(
new EagleInternalException(
s"Provider $providerName future succeeded before completed stream callback for " +
s"${LlmResolver.diagnosticSummary(llmRequest)}"
LlmResolver.incompleteStreamError(
providerName,
LlmResolver.diagnosticSummary(llmRequest)
)
)
}
@@ -12,7 +12,9 @@ scala_test(
name = "gemini_service_impl_test",
srcs = ["GeminiServiceImplTest.scala"],
deps = [
"//src/main/scala/net/eagle0/common/llm_integration:external_text_generation_error",
"//src/main/scala/net/eagle0/common/llm_integration:external_text_generation_service_impl",
"//src/main/scala/net/eagle0/common/llm_integration:gemini_service_impl",
"//src/main/scala/net/eagle0/common/llm_integration:streaming_text_results",
],
)
@@ -1,10 +1,27 @@
package net.eagle0.common.llm_integration
import java.util.function.Consumer
import scala.collection.mutable.ArrayBuffer
import org.scalatest.flatspec.AnyFlatSpec
import org.scalatest.matchers.should.Matchers
class GeminiServiceImplTest extends AnyFlatSpec with Matchers {
private def parse(responseJson: String): Vector[StreamingTextResults] = {
val results = ArrayBuffer.empty[StreamingTextResults]
new GeminiServiceImpl()
.stringConsumer(new Consumer[StreamingTextResults] {
override def accept(result: StreamingTextResults): Unit = results += result
})
.accept(responseJson)
results.toVector
}
private def parseFailure(responseJson: String): ExternalTextGenerationStreamFailure =
the[ExternalTextGenerationStreamFailure] thrownBy parse(responseJson)
"fallbackModelNames" should "try the configured model before fallback Gemini models" in {
GeminiServiceImpl.fallbackModelNames("gemini-2.5-flash-lite") shouldBe Vector(
"gemini-2.5-flash-lite",
@@ -28,4 +45,80 @@ class GeminiServiceImplTest extends AnyFlatSpec with Matchers {
.fallbackModelNames(GeminiServiceImpl.defaultModel)
.count(_ == GeminiServiceImpl.defaultModel) shouldBe 1
}
"stringConsumer" should "complete successful Gemini finish reasons" in
Vector("STOP", "MAX_TOKENS").foreach { finishReason =>
val results = parse(
s"""{
| "candidates": [{
| "content": {"parts": [{"text": "generated text"}]},
| "finishReason": "$finishReason"
| }],
| "usageMetadata": {"promptTokenCount": 11, "candidatesTokenCount": 7}
|}""".stripMargin
)
results.map(result => (result.value, result.completed)) shouldBe Vector("generated text" -> true)
results.head.tokenUsage shouldBe Some(TokenUsage(inputTokens = 11, outputTokens = 7))
}
it should "leave a candidate without a finish reason in progress" in {
val results = parse(
"""{
| "candidates": [{
| "content": {"parts": [{"text": "partial text"}]}
| }]
|}""".stripMargin
)
results.map(result => (result.value, result.completed)) shouldBe Vector("partial text" -> false)
}
it should "fail every non-success Gemini finish reason with its diagnostic details" in
Vector(
"SAFETY",
"RECITATION",
"LANGUAGE",
"OTHER",
"BLOCKLIST",
"PROHIBITED_CONTENT",
"SPII",
"MALFORMED_FUNCTION_CALL",
"IMAGE_SAFETY",
"FINISH_REASON_UNSPECIFIED"
).foreach { finishReason =>
val failure = parseFailure(
s"""{
| "candidates": [{
| "finishReason": "$finishReason",
| "finishMessage": "terminal details"
| }]
|}""".stripMargin
)
failure.getMessage should include(s"finishReason=$finishReason")
failure.getMessage should include("message=terminal details")
}
it should "fail a prompt-level block without candidates" in {
val failure = parseFailure(
"""{
| "promptFeedback": {
| "blockReason": "PROHIBITED_CONTENT",
| "blockReasonMessage": "prompt details"
| }
|}""".stripMargin
)
failure.getMessage should include("blockReason=PROHIBITED_CONTENT")
failure.getMessage should include("message=prompt details")
}
it should "fail an API error response" in {
val failure = parseFailure(
"""{"error": {"code": 503, "message": "provider unavailable"}}"""
)
failure.getMessage shouldBe "Gemini API error: provider unavailable"
}
}
@@ -368,6 +368,7 @@ scala_test(
timeout = "short",
srcs = ["LlmResolverTest.scala"],
deps = [
"//src/main/scala/net/eagle0/common/llm_integration:external_text_generation_error",
"//src/main/scala/net/eagle0/eagle:eagle_pkg",
"//src/main/scala/net/eagle0/eagle/client_text:text_generation_result",
"//src/main/scala/net/eagle0/eagle/model/action_result/generated_text_request",
@@ -200,6 +200,17 @@ class LlmResolverTest extends AnyFlatSpec with MockFactory {
it should "accept a completed stream that produced text" in
LlmResolver.completedWithoutText(streamedChars = 1, partialCompletion = None).shouldBe(false)
it should "classify a stream closed without a terminal callback as a provider failure" in {
val error = LlmResolver.incompleteStreamError(
providerName = "gemini/gemini-2.5-flash-lite",
requestSummary = "id=battalion_10_backstory"
)
error.code.shouldBe(502)
error.msg.should(include("gemini/gemini-2.5-flash-lite"))
error.msg.should(include("id=battalion_10_backstory"))
}
it should "sort success results that are dependencies first" in {
val sortedRequests =
llmResolver.sortedLlmRequestsWithPrompts(unsortedRequests)