fix(gateway): track oversized terminal response SSE events

Author: masterkain

lib/codex_pooler/gateway/runtime/streaming/downstream_stream.ex

@@ -29,6 +29,7 @@ defmodule CodexPooler.Gateway.Runtime.Streaming.DownstreamStream do
         |> Map.put(:public_openai_responses_data_seen?, false)
         |> Map.put(:public_openai_responses_terminal_seen?, false)
         |> Map.put(:public_openai_responses_response_id, nil)
+        |> Map.put(:public_openai_responses_terminal_kind, nil)
 
       true ->
         state
@@ -79,6 +80,13 @@ defmodule CodexPooler.Gateway.Runtime.Streaming.DownstreamStream do
 
   def keepalive_allowed?(_state), do: true
 
+  @spec terminal_outcome(state()) :: :completed | :incomplete | :failed | nil
+  def terminal_outcome(%{public_openai_responses_terminal_kind: kind})
+      when kind in [:completed, :incomplete, :failed],
+      do: kind
+
+  def terminal_outcome(_state), do: nil
+
   @spec synthetic_terminal_failure(state(), term()) :: {binary() | nil, state()}
   def synthetic_terminal_failure(
         %{
@@ -118,13 +126,26 @@ defmodule CodexPooler.Gateway.Runtime.Streaming.DownstreamStream do
     state =
       state
       |> Map.put(:public_openai_responses, stream_state)
+      |> maybe_mark_passthrough_terminal(stream_state)
       |> track_public_openai_responses_output(data)
 
     {data, state}
   end
 
   defp normalize_public_openai_responses_stream_data(data, state), do: {data, state}
 
+  defp maybe_mark_passthrough_terminal(state, stream_state) do
+    case StreamProtocol.public_openai_responses_passthrough_terminal_kind(stream_state) do
+      kind when kind in [:completed, :incomplete, :failed] ->
+        state
+        |> Map.put(:public_openai_responses_terminal_seen?, true)
+        |> Map.put(:public_openai_responses_terminal_kind, kind)
+
+      _kind ->
+        state
+    end
+  end
+
   defp track_public_openai_responses_output(state, data) when is_binary(data) do
     {blocks, _buffer} = StreamProtocol.complete_sse_blocks(data, bounded?: false)
 
@@ -134,7 +155,9 @@ defmodule CodexPooler.Gateway.Runtime.Streaming.DownstreamStream do
   end
 
   defp track_public_openai_responses_block(state, "data: [DONE]") do
-    %{state | public_openai_responses_terminal_seen?: true}
+    state
+    |> Map.put(:public_openai_responses_terminal_seen?, true)
+    |> Map.put(:public_openai_responses_terminal_kind, :completed)
   end
 
   defp track_public_openai_responses_block(state, block) do
@@ -249,10 +272,14 @@ defmodule CodexPooler.Gateway.Runtime.Streaming.DownstreamStream do
   end
 
   defp maybe_mark_public_openai_responses_terminal_seen(state, event) do
-    if StreamProtocol.terminal_outcome_event(event) do
-      %{state | public_openai_responses_terminal_seen?: true}
-    else
-      state
+    case StreamProtocol.terminal_outcome_event(event) do
+      {:ok, %{kind: kind}} when kind in [:completed, :incomplete, :failed] ->
+        state
+        |> Map.put(:public_openai_responses_terminal_seen?, true)
+        |> Map.put(:public_openai_responses_terminal_kind, kind)
+
+      _outcome ->
+        state
     end
   end
 

lib/codex_pooler/gateway/runtime/streaming/stream_dispatch.ex

@@ -218,20 +218,39 @@ defmodule CodexPooler.Gateway.Runtime.Streaming.StreamDispatch do
 
   defp http_stream_terminal_failure_writer do
     fn state, reason ->
-      write_public_openai_responses_terminal_failure(state, reason)
+      case DownstreamStream.terminal_outcome(state) do
+        terminal when terminal in [:completed, :incomplete] ->
+          {:success, state, ""}
+
+        _missing_or_failed_terminal ->
+          write_public_openai_responses_terminal_failure(state, reason)
+      end
     end
   end
 
   defp http_stream_terminal_success_hook do
     fn state ->
-      case write_public_openai_responses_terminal_failure(state, :upstream_stream_interrupted) do
-        {:ok, state, ""} -> {:ok, state, ""}
-        {:ok, state, data} -> {:failure, state, data, :upstream_stream_interrupted}
-        {:error, _reason} -> {:failure, state, "", :upstream_stream_interrupted}
+      case DownstreamStream.terminal_outcome(state) do
+        :failed ->
+          {:failure, state, "", :upstream_stream_interrupted}
+
+        terminal when terminal in [:completed, :incomplete] ->
+          {:ok, state, ""}
+
+        _missing_terminal ->
+          missing_public_openai_responses_terminal_result(state)
       end
     end
   end
 
+  defp missing_public_openai_responses_terminal_result(state) do
+    case write_public_openai_responses_terminal_failure(state, :upstream_stream_interrupted) do
+      {:ok, state, ""} -> {:ok, state, ""}
+      {:ok, state, data} -> {:failure, state, data, :upstream_stream_interrupted}
+      {:error, _reason} -> {:failure, state, "", :upstream_stream_interrupted}
+    end
+  end
+
   defp write_public_openai_responses_terminal_failure(state, reason) do
     case DownstreamStream.synthetic_terminal_failure(state, reason) do
       {nil, state} ->

lib/codex_pooler/gateway/transports/streaming/stream_protocol.ex

@@ -64,7 +64,10 @@ defmodule CodexPooler.Gateway.Transports.Streaming.StreamProtocol do
           required(:buffer) => binary(),
           required(:created?) => boolean(),
           required(:text_delta?) => boolean(),
-          required(:passthrough?) => boolean()
+          required(:passthrough?) => boolean(),
+          required(:passthrough_terminal) => PublicResponses.passthrough_terminal_state() | nil,
+          required(:passthrough_terminal_kind) => atom() | nil,
+          required(:passthrough_terminal_seen?) => boolean()
         }
   @type websocket_frame_headers :: %{optional(String.t()) => String.t()}
 
@@ -122,6 +125,11 @@ defmodule CodexPooler.Gateway.Transports.Streaming.StreamProtocol do
   def normalize_public_openai_responses_sse_data(data, state),
     do: PublicResponses.normalize_data(data, state)
 
+  @spec public_openai_responses_passthrough_terminal_kind(public_openai_responses_stream_state()) ::
+          atom() | nil
+  def public_openai_responses_passthrough_terminal_kind(state),
+    do: PublicResponses.passthrough_terminal_kind(state)
+
   @spec synthetic_public_openai_responses_failure_sse(String.t() | nil, term()) :: binary()
   def synthetic_public_openai_responses_failure_sse(response_id, _reason) do
     error = %{
@@ -399,14 +407,22 @@ defmodule CodexPooler.Gateway.Transports.Streaming.StreamProtocol do
 
   @spec sse_field(binary(), binary()) :: binary() | nil
   def sse_field(block, name) do
-    prefix = name <> ": "
+    prefix = name <> ":"
 
     block
     |> String.split("\n")
     |> Enum.map(&String.trim/1)
-    |> Enum.find_value(fn line ->
-      if String.starts_with?(line, prefix), do: String.replace_prefix(line, prefix, "")
+    |> Enum.flat_map(fn line ->
+      if String.starts_with?(line, prefix) do
+        [line |> String.replace_prefix(prefix, "") |> String.trim_leading()]
+      else
+        []
+      end
     end)
+    |> case do
+      [] -> nil
+      values -> Enum.join(values, "\n")
+    end
   end
 
   @spec decode_sse_data(term()) :: map()

lib/codex_pooler/gateway/transports/streaming/stream_protocol/public_responses.ex

@@ -5,15 +5,36 @@ defmodule CodexPooler.Gateway.Transports.Streaming.StreamProtocol.PublicResponse
   alias CodexPooler.Gateway.Runtime.Streaming.BufferTelemetry
   alias CodexPooler.Gateway.Transports.Streaming.StreamProtocol
 
+  @type passthrough_terminal_state :: %{
+          required(:event_type) => String.t(),
+          required(:json_started?) => boolean(),
+          required(:depth) => non_neg_integer(),
+          required(:in_string?) => boolean(),
+          required(:escaped?) => boolean(),
+          required(:complete?) => boolean()
+        }
   @type state :: %{
           required(:buffer) => binary(),
           required(:created?) => boolean(),
           required(:text_delta?) => boolean(),
-          required(:passthrough?) => boolean()
+          required(:passthrough?) => boolean(),
+          required(:passthrough_terminal) => passthrough_terminal_state() | nil,
+          required(:passthrough_terminal_kind) => atom() | nil,
+          required(:passthrough_terminal_seen?) => boolean()
         }
 
   @spec new_state() :: state()
-  def new_state, do: %{buffer: "", created?: false, text_delta?: false, passthrough?: false}
+  def new_state do
+    %{
+      buffer: "",
+      created?: false,
+      text_delta?: false,
+      passthrough?: false,
+      passthrough_terminal: nil,
+      passthrough_terminal_kind: nil,
+      passthrough_terminal_seen?: false
+    }
+  end
 
   @terminal_buffer_markers [
     "data: [DONE]",
@@ -40,12 +61,12 @@ defmodule CodexPooler.Gateway.Transports.Streaming.StreamProtocol.PublicResponse
 
       blocks == [] and StreamProtocol.oversized_incomplete_sse_block?(buffer) ->
         record_oversized_incomplete(byte_size(buffered_data))
-        {buffered_data, %{state | buffer: "", passthrough?: true}}
+        {buffered_data, enter_passthrough(buffered_data, state)}
 
       StreamProtocol.oversized_incomplete_sse_block?(buffer) ->
         record_oversized_incomplete(byte_size(buffered_data))
         {iodata, state} = normalize_complete_blocks(blocks, state)
-        {[iodata, buffer] |> IO.iodata_to_binary(), %{state | buffer: "", passthrough?: true}}
+        {[iodata, buffer] |> IO.iodata_to_binary(), enter_passthrough(buffer, state)}
 
       true ->
         normalize_blocks(blocks, buffer, state)
@@ -61,13 +82,165 @@ defmodule CodexPooler.Gateway.Transports.Streaming.StreamProtocol.PublicResponse
         passthrough = binary_part(data, 0, passthrough_size)
         rest = binary_part(data, passthrough_size, byte_size(data) - passthrough_size)
 
-        state = %{state | passthrough?: false, buffer: ""}
+        state =
+          state
+          |> track_passthrough_terminal_data(passthrough)
+          |> then(&%{&1 | passthrough?: false, buffer: "", passthrough_terminal: nil})
+
         {normalized_rest, state} = normalize_data(rest, state)
 
         {[passthrough, normalized_rest] |> IO.iodata_to_binary(), state}
 
       nil ->
-        {data, state}
+        {data, track_passthrough_terminal_data(state, data)}
+    end
+  end
+
+  @spec passthrough_terminal_kind(state()) :: atom() | nil
+  def passthrough_terminal_kind(%{passthrough_terminal_kind: kind}) when is_atom(kind), do: kind
+  def passthrough_terminal_kind(_state), do: nil
+
+  defp enter_passthrough(data, state) do
+    data
+    |> passthrough_terminal_event_type()
+    |> case do
+      event_type when is_binary(event_type) ->
+        state
+        |> Map.put(:passthrough_terminal, new_passthrough_terminal(event_type))
+        |> track_passthrough_terminal_data(data)
+
+      nil ->
+        %{state | passthrough_terminal: nil}
+    end
+    |> then(&%{&1 | buffer: "", passthrough?: true})
+  end
+
+  defp new_passthrough_terminal(event_type) do
+    %{
+      event_type: event_type,
+      json_started?: false,
+      depth: 0,
+      in_string?: false,
+      escaped?: false,
+      complete?: false
+    }
+  end
+
+  defp passthrough_terminal_event_type(data) do
+    {event_type, decoded} = stream_block_event(data)
+    type = event_type || decoded_string(decoded, "type") || terminal_marker_type(data)
+
+    if terminal_event?(type), do: type
+  end
+
+  defp terminal_marker_type(data) when is_binary(data) do
+    cond do
+      terminal_marker?(data, "response.completed") -> "response.completed"
+      terminal_marker?(data, "response.failed") -> "response.failed"
+      terminal_marker?(data, "response.incomplete") -> "response.incomplete"
+      terminal_marker?(data, "error") -> "error"
+      true -> nil
+    end
+  end
+
+  defp terminal_marker?(data, type) do
+    String.contains?(data, "event: #{type}") or
+      String.contains?(data, "event:#{type}") or
+      String.contains?(data, ~s("type":"#{type}")) or
+      String.contains?(data, ~s("type": "#{type}"))
+  end
+
+  defp track_passthrough_terminal_data(
+         %{passthrough_terminal: %{complete?: false} = terminal} = state,
+         data
+       )
+       when is_binary(data) do
+    terminal = scan_passthrough_terminal_json(data, terminal)
+
+    state
+    |> Map.put(:passthrough_terminal, terminal)
+    |> maybe_put_passthrough_terminal_kind(terminal)
+    |> Map.put(
+      :passthrough_terminal_seen?,
+      state.passthrough_terminal_seen? or terminal.complete?
+    )
+  end
+
+  defp track_passthrough_terminal_data(state, _data), do: state
+
+  defp maybe_put_passthrough_terminal_kind(state, %{complete?: true, event_type: event_type}) do
+    Map.put(state, :passthrough_terminal_kind, terminal_kind(event_type))
+  end
+
+  defp maybe_put_passthrough_terminal_kind(state, _terminal), do: state
+
+  defp terminal_kind("response.completed"), do: :completed
+  defp terminal_kind("response.incomplete"), do: :incomplete
+  defp terminal_kind(_event_type), do: :failed
+
+  defp scan_passthrough_terminal_json(data, terminal),
+    do: scan_passthrough_terminal_json(data, 0, terminal)
+
+  defp scan_passthrough_terminal_json(data, offset, terminal)
+       when offset >= byte_size(data) or terminal.complete?,
+       do: terminal
+
+  defp scan_passthrough_terminal_json(data, offset, %{json_started?: false} = terminal) do
+    case :binary.at(data, offset) do
+      ?{ ->
+        scan_passthrough_terminal_json(data, offset + 1, %{
+          terminal
+          | json_started?: true,
+            depth: 1
+        })
+
+      _byte ->
+        scan_passthrough_terminal_json(data, offset + 1, terminal)
+    end
+  end
+
+  defp scan_passthrough_terminal_json(
+         data,
+         offset,
+         %{in_string?: true, escaped?: true} = terminal
+       ) do
+    scan_passthrough_terminal_json(data, offset + 1, %{terminal | escaped?: false})
+  end
+
+  defp scan_passthrough_terminal_json(
+         data,
+         offset,
+         %{in_string?: true, escaped?: false} = terminal
+       ) do
+    case :binary.at(data, offset) do
+      ?\\ -> scan_passthrough_terminal_json(data, offset + 1, %{terminal | escaped?: true})
+      ?" -> scan_passthrough_terminal_json(data, offset + 1, %{terminal | in_string?: false})
+      _byte -> scan_passthrough_terminal_json(data, offset + 1, terminal)
+    end
+  end
+
+  defp scan_passthrough_terminal_json(data, offset, terminal) do
+    case :binary.at(data, offset) do
+      ?" ->
+        scan_passthrough_terminal_json(data, offset + 1, %{terminal | in_string?: true})
+
+      ?{ ->
+        scan_passthrough_terminal_json(data, offset + 1, %{
+          terminal
+          | depth: terminal.depth + 1
+        })
+
+      ?} when terminal.depth <= 1 ->
+        %{terminal | depth: 0, complete?: true}
+
+      ?} ->
+        scan_passthrough_terminal_json(data, offset + 1, %{
+          terminal
+          | depth: terminal.depth - 1
+        })
+
+      _byte ->
+        scan_passthrough_terminal_json(data, offset + 1, terminal)
     end
   end