seewe/sub2api
1
0
Fork 0
mirror of https://github.com/Wei-Shaw/sub2api.git synced 2026-10-14 12:02:27 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix(antigravity): keep compat streams alive before first content

Browse Source
This commit is contained in:
wucm667
2026-09-28 12:15:27 +08:00
parent 9a62841fd1
commit bd48b2b2d2
2 changed files with 335 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/internal/service/antigravity_gateway_compat_stream.go
+94 -15
View File
@@ -107,16 +107,24 @@ type antigravityCompatScanEvent struct {
}
type antigravityCompatStreamSession struct {
processor *antigravity.StreamingProcessor
adapter antigravityCompatStreamAdapter
writer *antigravityClientWriter
usage *ClaudeUsage
pendingEvents []apicompat.AnthropicStreamEvent
firstTokenMs *int
startTime time.Time
meaningfulData bool
processor *antigravity.StreamingProcessor
adapter antigravityCompatStreamAdapter
writer *antigravityClientWriter
usage *ClaudeUsage
pendingEvents []apicompat.AnthropicStreamEvent
firstTokenMs *int
startTime time.Time
meaningfulData bool
preContentKeepaliveSent bool
preContentKeepaliveInterval time.Duration
}
const (
antigravityCompatPreContentKeepaliveInterval = 15 * time.Second
// Allow slow first tokens, but do not let comment-only streams reset the idle timer forever.
antigravityCompatPreContentMaxWait = 2 * time.Minute
)
func newAntigravityCompatStreamSession(
model string,
startTime time.Time,
@@ -124,11 +132,12 @@ func newAntigravityCompatStreamSession(
writer *antigravityClientWriter,
) *antigravityCompatStreamSession {
return &antigravityCompatStreamSession{
processor: antigravity.NewStreamingProcessor(model),
adapter: adapter,
writer: writer,
usage: &ClaudeUsage{},
startTime: startTime,
processor: antigravity.NewStreamingProcessor(model),
adapter: adapter,
writer: writer,
usage: &ClaudeUsage{},
startTime: startTime,
preContentKeepaliveInterval: antigravityCompatPreContentKeepaliveInterval,
}
}
@@ -144,11 +153,26 @@ func (s *antigravityCompatStreamSession) hasMeaningfulData() bool {
return s.meaningfulData || s.processor.HasContent()
}
func (s *antigravityCompatStreamSession) writePreContentKeepalive(now time.Time) {
// Intentional tradeoff: this commits HTTP 200 after 15s, so later upstream
// failures cannot fail over; report them as SSE errors to the client instead.
if s.hasMeaningfulData() || s.writer.Disconnected() || now.Sub(s.startTime) < s.preContentKeepaliveInterval {
return
}
if s.writer.Write([]byte(": ping\n\n")) {
s.preContentKeepaliveSent = true
}
}
func (s *antigravityCompatStreamSession) finish() (*antigravityStreamResult, error) {
finalEvents, usage := s.processor.Finish()
mergeAntigravityCompatUsage(s.usage, usage)
s.consumeClaudeEvents(finalEvents)
if !s.hasMeaningfulData() && !s.writer.Disconnected() {
if s.preContentKeepaliveSent {
s.adapter.WriteError(s.writer, "empty_stream")
return s.result(false), errors.New("empty Antigravity compatibility stream after keepalive")
}
return nil, antigravityCompatEmptyStreamError()
}
s.adapter.Finalize(s.writer)
@@ -257,6 +281,13 @@ func (s *AntigravityGatewayService) handleAntigravityCompatStream(
originalModel string,
adapter antigravityCompatStreamAdapter,
prefix string,
) (*antigravityStreamResult, error) {
return s.handleAntigravityCompatStreamWithKeepaliveInterval(c, resp, startTime, originalModel, adapter, prefix, antigravityCompatPreContentKeepaliveInterval, antigravityCompatPreContentMaxWait)
}
func (s *AntigravityGatewayService) handleAntigravityCompatStreamWithKeepaliveInterval(
c *gin.Context, resp *http.Response, startTime time.Time, originalModel string,
adapter antigravityCompatStreamAdapter, prefix string, interval, maxPreContentWait time.Duration,
) (*antigravityStreamResult, error) {
flusher, ok := c.Writer.(http.Flusher)
if !ok {
@@ -272,6 +303,7 @@ func (s *AntigravityGatewayService) handleAntigravityCompatStream(
c.Status(http.StatusOK)
}
session := newAntigravityCompatStreamSession(originalModel, startTime, adapter, writer)
session.preContentKeepaliveInterval = interval
events, stopScanner, maxLineSize := s.startAntigravityCompatScanner(resp.Body)
defer stopScanner()
@@ -284,13 +316,36 @@ func (s *AntigravityGatewayService) handleAntigravityCompatStream(
if keepaliveTicker != nil {
defer keepaliveTicker.Stop()
}
preContentDelay := time.Until(startTime.Add(interval))
if preContentDelay < 0 {
preContentDelay = 0
}
preContentTimer := time.NewTimer(preContentDelay)
defer preContentTimer.Stop()
preContentDeadline := startTime.Add(maxPreContentWait)
deadlineDelay := time.Until(preContentDeadline)
if deadlineDelay < 0 {
deadlineDelay = 0
}
deadlineTimer := time.NewTimer(deadlineDelay)
defer deadlineTimer.Stop()
preContentTimeout := func() (*antigravityStreamResult, error) {
if session.preContentKeepaliveSent {
writeAntigravityCompatStreamError(c, adapter, writer, "stream_timeout")
return session.collectResult(false), fmt.Errorf("pre-content stream timeout")
}
return nil, antigravityCompatEmptyStreamError()
}
for {
select {
case event, open := <-events:
if !session.hasMeaningfulData() && !writer.Disconnected() && !time.Now().Before(preContentDeadline) {
return preContentTimeout()
}
if !open {
if !session.hasMeaningfulData() && !writer.Disconnected() {
return nil, antigravityCompatEmptyStreamError()
return handleAntigravityCompatEmptyStream(c, session)
}
return session.finish()
}
@@ -306,7 +361,7 @@ func (s *AntigravityGatewayService) handleAntigravityCompatStream(
return session.collectResult(true), nil
}
if !session.hasMeaningfulData() {
return nil, antigravityCompatEmptyStreamError()
return preContentTimeout()
}
logger.LegacyPrintf("service.antigravity_gateway", "Stream data interval timeout (%s)", prefix)
writeAntigravityCompatStreamError(c, adapter, writer, "stream_timeout")
@@ -316,10 +371,30 @@ func (s *AntigravityGatewayService) handleAntigravityCompatStream(
if session.hasMeaningfulData() && !writer.Disconnected() {
writer.Write([]byte(": ping\n\n"))
}
case now := <-preContentTimer.C:
if !session.hasMeaningfulData() && !writer.Disconnected() && !now.Before(preContentDeadline) {
return preContentTimeout()
}
session.writePreContentKeepalive(now)
if !session.hasMeaningfulData() && !writer.Disconnected() {
preContentTimer.Reset(interval)
}
case <-deadlineTimer.C:
if !session.hasMeaningfulData() && !writer.Disconnected() {
return preContentTimeout()
}
}
}
}
func handleAntigravityCompatEmptyStream(c *gin.Context, session *antigravityCompatStreamSession) (*antigravityStreamResult, error) {
if session.preContentKeepaliveSent {
writeAntigravityCompatStreamError(c, session.adapter, session.writer, "empty_stream")
return session.collectResult(false), errors.New("empty Antigravity compatibility stream after keepalive")
}
return nil, antigravityCompatEmptyStreamError()
}
func (s *AntigravityGatewayService) startAntigravityCompatScanner(
body io.Reader,
) (<-chan antigravityCompatScanEvent, func(), int) {
@@ -404,6 +479,10 @@ func (s *AntigravityGatewayService) handleAntigravityCompatReadError(
prefix string,
) (*antigravityStreamResult, error) {
if !session.hasMeaningfulData() && !session.writer.Disconnected() {
if session.preContentKeepaliveSent {
writeAntigravityCompatStreamError(c, session.adapter, session.writer, "stream_read_error")
return session.collectResult(false), fmt.Errorf("stream read error: %w", err)
}
return nil, antigravityCompatEmptyStreamError()
}
if disconnect, handled := handleStreamReadError(err, session.writer.Disconnected(), prefix); handled {
backend/internal/service/antigravity_gateway_compat_test.go
+241
View File
@@ -28,6 +28,20 @@ type antigravityCompatErrorReader struct {
err error
}
type antigravityCompatNotifyingWriter struct {
gin.ResponseWriter
wrote chan struct{}
}
func (w *antigravityCompatNotifyingWriter) Write(data []byte) (int, error) {
n, err := w.ResponseWriter.Write(data)
select {
case w.wrote <- struct{}{}:
default:
}
return n, err
}
func (r *antigravityCompatErrorReader) Read(p []byte) (int, error) {
if r.off < len(r.data) {
n := copy(p, r.data[r.off:])
@@ -786,3 +800,230 @@ func TestAntigravityCompatKeepaliveAfterFirstEvent(t *testing.T) {
require.Contains(t, recorder.Header().Get("Content-Type"), "text/event-stream")
require.NoError(t, reader.Close())
}
func TestAntigravityCompatPreContentKeepalive(t *testing.T) {
gin.SetMode(gin.TestMode)
for _, tt := range []struct {
name string
adapter antigravityCompatStreamAdapter
want string
}{
{"chat completions", newAntigravityChatStreamAdapter("gemini-3.1-pro", false), "data: "},
{"responses", newAntigravityResponsesStreamAdapter("gemini-3.1-pro"), "event: "},
} {
t.Run(tt.name, func(t *testing.T) {
c, recorder := newAntigravityCompatContext(http.MethodPost, "/", nil)
writer := newAntigravityClientWriter(c.Writer, c.Writer, "test")
writer.beforeFirstWrite = func() { c.Header("Content-Type", "text/event-stream") }
start := time.Now().Add(-20 * time.Second)
session := newAntigravityCompatStreamSession("gemini-3.1-pro", start, tt.adapter, writer)
session.writePreContentKeepalive(start.Add(14 * time.Second))
require.Empty(t, recorder.Body.String())
session.consumeClaudeData("message_start", `{"type":"message_start","message":{"id":"msg_1","type":"message","role":"assistant"}}`)
require.False(t, session.hasMeaningfulData())
session.writePreContentKeepalive(start.Add(14 * time.Second))
require.Empty(t, recorder.Body.String())
session.writePreContentKeepalive(start.Add(15 * time.Second))
require.Equal(t, ": ping\n\n", recorder.Body.String())
require.Nil(t, session.firstTokenMs)
session.consumeClaudeData("content_block_delta", `{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"hello"}}`)
require.True(t, session.hasMeaningfulData())
require.NotNil(t, session.firstTokenMs)
require.GreaterOrEqual(t, *session.firstTokenMs, 20000)
require.Contains(t, recorder.Body.String(), tt.want)
require.Greater(t, strings.Index(recorder.Body.String(), tt.want), strings.Index(recorder.Body.String(), ": ping"))
before := recorder.Body.String()
session.writePreContentKeepalive(start.Add(30 * time.Second))
require.Equal(t, before, recorder.Body.String())
})
}
}
func TestAntigravityCompatHandlerPreContentKeepalive(t *testing.T) {
gin.SetMode(gin.TestMode)
for _, tt := range []struct {
name string
adapter func() antigravityCompatStreamAdapter
line string
want string
}{
{"silent chat", func() antigravityCompatStreamAdapter { return newAntigravityChatStreamAdapter("gemini-3.1-pro", false) }, "", `"upstream_error"`},
{"signature-only responses", func() antigravityCompatStreamAdapter { return newAntigravityResponsesStreamAdapter("gemini-3.1-pro") }, `data: {"response":{"candidates":[{"content":{"parts":[{"thoughtSignature":"sig","text":""}]}}]}}` + "\n\n", "event: error"},
} {
t.Run(tt.name, func(t *testing.T) {
svc := newAntigravityCompatService(config.GatewayConfig{MaxLineSize: defaultMaxLineSize, StreamDataIntervalTimeout: 30}, nil)
c, recorder := newAntigravityCompatContext(http.MethodPost, "/", nil)
notifier := &antigravityCompatNotifyingWriter{ResponseWriter: c.Writer, wrote: make(chan struct{}, 1)}
c.Writer = notifier
reader, pipeWriter := io.Pipe()
defer func() { _ = reader.Close() }()
resp := &http.Response{StatusCode: http.StatusOK, Body: reader}
done := make(chan error, 1)
go func() {
_, err := svc.handleAntigravityCompatStream(c, resp, time.Now().Add(-antigravityCompatPreContentKeepaliveInterval+200*time.Millisecond), "gemini-3.1-pro", tt.adapter(), "test")
done <- err
}()
if tt.line != "" {
_, err := io.WriteString(pipeWriter, tt.line)
require.NoError(t, err)
}
select {
case <-notifier.wrote:
case <-time.After(2 * time.Second):
_ = pipeWriter.Close()
t.Fatal("no pre-content keepalive before read timeout")
}
require.Equal(t, ": ping\n\n", recorder.Body.String())
require.NoError(t, pipeWriter.Close())
require.Error(t, <-done)
require.Contains(t, recorder.Body.String(), tt.want)
require.True(t, IsResponseCommitted(c))
})
}
}
func TestAntigravityCompatHandlerRepeatsPreContentKeepalive(t *testing.T) {
svc := newAntigravityCompatService(config.GatewayConfig{
MaxLineSize: defaultMaxLineSize, StreamDataIntervalTimeout: 30, StreamKeepaliveInterval: 0,
}, nil)
c, recorder := newAntigravityCompatContext(http.MethodPost, "/", nil)
reader, pipeWriter := io.Pipe()
defer func() { _ = reader.Close() }()
resp := &http.Response{StatusCode: http.StatusOK, Body: reader}
done := make(chan error, 1)
go func() {
_, err := svc.handleAntigravityCompatStreamWithKeepaliveInterval(
c, resp, time.Now().Add(-15*time.Millisecond), "gemini-3.1-pro",
newAntigravityResponsesStreamAdapter("gemini-3.1-pro"), "test", 15*time.Millisecond, time.Second,
)
done <- err
}()
time.Sleep(55 * time.Millisecond)
require.NoError(t, pipeWriter.Close())
require.Error(t, <-done)
require.GreaterOrEqual(t, strings.Count(recorder.Body.String(), ": ping\n\n"), 3)
require.Contains(t, recorder.Body.String(), "event: error")
require.True(t, IsResponseCommitted(c))
}
func TestAntigravityCompatHandlerPreContentDeadlineWithCommentOnlyStream(t *testing.T) {
gin.SetMode(gin.TestMode)
svc := newAntigravityCompatService(config.GatewayConfig{MaxLineSize: defaultMaxLineSize, StreamDataIntervalTimeout: 1}, nil)
c, recorder := newAntigravityCompatContext(http.MethodPost, "/v1/responses", nil)
reader, pipeWriter := io.Pipe()
defer func() { _ = reader.Close() }()
resp := &http.Response{StatusCode: http.StatusOK, Body: reader}
done := make(chan error, 1)
go func() {
_, err := svc.handleAntigravityCompatStreamWithKeepaliveInterval(c, resp, time.Now(), "gemini-3.1-pro",
newAntigravityResponsesStreamAdapter("gemini-3.1-pro"), "test", 10*time.Millisecond, 100*time.Millisecond)
done <- err
}()
stop := make(chan struct{})
defer close(stop)
go func() {
ticker := time.NewTicker(5 * time.Millisecond)
defer ticker.Stop()
for {
select {
case <-stop:
return
case <-ticker.C:
if _, err := io.WriteString(pipeWriter, ": upstream ping\n\n"); err != nil {
return
}
}
}
}()
select {
case err := <-done:
require.ErrorContains(t, err, "pre-content stream timeout")
case <-time.After(time.Second):
t.Fatal("comment-only stream exceeded pre-content deadline")
}
require.Contains(t, recorder.Body.String(), ": ping\n\n")
require.Contains(t, recorder.Body.String(), "event: error")
require.True(t, IsResponseCommitted(c))
_ = pipeWriter.Close()
}
func TestAntigravityCompatExpiredPreContentDeadlineDoesNotCommit(t *testing.T) {
gin.SetMode(gin.TestMode)
svc := newAntigravityCompatService(config.GatewayConfig{MaxLineSize: defaultMaxLineSize}, nil)
c, recorder := newAntigravityCompatContext(http.MethodPost, "/v1/responses", nil)
reader, pipeWriter := io.Pipe()
defer func() { _ = reader.Close() }()
defer func() { _ = pipeWriter.Close() }()
resp := &http.Response{StatusCode: http.StatusOK, Body: reader}
result, err := svc.handleAntigravityCompatStreamWithKeepaliveInterval(c, resp,
time.Now().Add(-time.Minute), "gemini-3.1-pro",
newAntigravityResponsesStreamAdapter("gemini-3.1-pro"), "test", 10*time.Millisecond, 20*time.Millisecond)
require.Nil(t, result)
var failoverErr *UpstreamFailoverError
require.ErrorAs(t, err, &failoverErr)
require.Empty(t, recorder.Body.String())
require.False(t, IsResponseCommitted(c))
}
func TestAntigravityCompatHandlerErrorsAfterPreContentPing(t *testing.T) {
gin.SetMode(gin.TestMode)
for _, tc := range []struct {
name string
wait time.Duration
want string
}{
{"read error", time.Second, "stream_read_error"},
{"timeout", 25 * time.Millisecond, "stream_timeout"},
} {
t.Run(tc.name, func(t *testing.T) {
svc := newAntigravityCompatService(config.GatewayConfig{MaxLineSize: defaultMaxLineSize}, nil)
c, recorder := newAntigravityCompatContext(http.MethodPost, "/v1/responses", nil)
reader, pipeWriter := io.Pipe()
defer func() { _ = reader.Close() }()
defer func() { _ = pipeWriter.Close() }()
if tc.name == "read error" {
go func() {
time.Sleep(15 * time.Millisecond)
_ = pipeWriter.CloseWithError(io.ErrUnexpectedEOF)
}()
}
resp := &http.Response{StatusCode: http.StatusOK, Body: reader}
_, err := svc.handleAntigravityCompatStreamWithKeepaliveInterval(c, resp,
time.Now().Add(-20*time.Millisecond), "gemini-3.1-pro",
newAntigravityResponsesStreamAdapter("gemini-3.1-pro"), "test", 10*time.Millisecond, tc.wait)
require.Error(t, err)
require.Contains(t, recorder.Body.String(), ": ping\n\n")
require.Contains(t, recorder.Body.String(), tc.want)
require.True(t, IsResponseCommitted(c))
})
}
}
func TestAntigravityCompatEmptyAfterKeepaliveReportsStreamError(t *testing.T) {
gin.SetMode(gin.TestMode)
for _, tt := range []struct {
name string
adapter antigravityCompatStreamAdapter
want string
}{
{"chat completions", newAntigravityChatStreamAdapter("gemini-3.1-pro", false), `"upstream_error"`},
{"responses", newAntigravityResponsesStreamAdapter("gemini-3.1-pro"), "event: error"},
} {
t.Run(tt.name, func(t *testing.T) {
c, recorder := newAntigravityCompatContext(http.MethodPost, "/", nil)
writer := newAntigravityClientWriter(c.Writer, c.Writer, "test")
writer.beforeFirstWrite = func() { c.Header("Content-Type", "text/event-stream") }
start := time.Now().Add(-20 * time.Second)
session := newAntigravityCompatStreamSession("gemini-3.1-pro", start, tt.adapter, writer)
session.consumeClaudeData("message_start", `{"type":"message_start","message":{"id":"msg_1","type":"message","role":"assistant"}}`)
session.writePreContentKeepalive(start.Add(15 * time.Second))
result, err := handleAntigravityCompatEmptyStream(c, session)
require.Error(t, err)
var failoverErr *UpstreamFailoverError
require.NotErrorAs(t, err, &failoverErr)
require.NotNil(t, result)
require.True(t, IsResponseCommitted(c))
require.Contains(t, recorder.Body.String(), tt.want)
})
}
}