From b963e13bf9ffd60928b1363bd0988b823db1e610 Mon Sep 17 00:00:00 2001 From: toki Date: Tue, 9 Jun 2026 04:13:27 +0900 Subject: [PATCH] =?UTF-8?q?fix(edge):=20run=20stream=20=EC=A2=85=EB=A3=8C?= =?UTF-8?q?=20=EC=8B=9C=20=EC=8B=A4=ED=8C=A8=20=EA=B2=BD=EB=A1=9C=EB=A5=BC?= =?UTF-8?q?=20=EA=B3=A0=EC=A0=95=ED=95=9C=EB=8B=A4?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit run stream 또는 node event 채널이 조기에 닫히는 경우에도 nil/closed 상태를 안전하게 처리해 무한 대기를 막고, A2A/OpenAI 경로가 일관되게 오류를 반환하도록 하여 예외 상태에서 서비스 장애 전파를 개선한다. --- apps/edge/internal/input/a2a/task_store.go | 14 +++++- .../internal/input/a2a/task_store_test.go | 20 +++++++++ apps/edge/internal/openai/run_result.go | 14 +++++- apps/edge/internal/openai/server_test.go | 43 +++++++++++++++++++ apps/edge/internal/openai/stream.go | 16 ++++++- apps/edge/internal/service/run_dispatch.go | 6 ++- .../service/run_dispatch_internal_test.go | 19 ++++++++ 7 files changed, 124 insertions(+), 8 deletions(-) create mode 100644 apps/edge/internal/service/run_dispatch_internal_test.go diff --git a/apps/edge/internal/input/a2a/task_store.go b/apps/edge/internal/input/a2a/task_store.go index 7d54664..c2ff657 100644 --- a/apps/edge/internal/input/a2a/task_store.go +++ b/apps/edge/internal/input/a2a/task_store.go @@ -99,13 +99,23 @@ func (s *TaskStore) drain(taskID string, handle edgeservice.RunResult) *Task { defer timeout.Stop() stream := handle.Stream() + if stream.Events == nil { + return s.setFailed(taskID, "run stream unavailable") + } for { select { - case nodeEvent := <-stream.NodeEvents: + case nodeEvent, ok := <-stream.NodeEvents: + if !ok { + stream.NodeEvents = nil + continue + } if edgeservice.IsNodeDisconnected(nodeEvent) { return s.setFailed(taskID, "node disconnected") } - case event := <-stream.Events: + case event, ok := <-stream.Events: + if !ok { + return s.setFailed(taskID, "run stream closed") + } if event == nil { continue } diff --git a/apps/edge/internal/input/a2a/task_store_test.go b/apps/edge/internal/input/a2a/task_store_test.go index b07bac5..81094ec 100644 --- a/apps/edge/internal/input/a2a/task_store_test.go +++ b/apps/edge/internal/input/a2a/task_store_test.go @@ -106,6 +106,26 @@ func TestTaskStoreMapsErrorToFailedTask(t *testing.T) { ) } +func TestTaskStoreMapsClosedRunStreamToFailedTask(t *testing.T) { + runID := "run-ts-closed" + events := make(chan *iop.RunEvent) + close(events) + handle := &edgeservice.RunHandle{ + RunDispatch: edgeservice.RunDispatch{RunID: runID, TimeoutSec: 60}, + RunStream: edgeservice.RunStream{ + Events: events, + NodeEvents: make(chan *iop.EdgeNodeEvent), + }, + } + + checkBlockingResult(t, handle, func(task map[string]any) { + status := task["status"].(map[string]any) + if status["state"] != "failed" { + t.Errorf("expected failed, got %q", status["state"]) + } + }) +} + func TestMessageSendBlockingReturnsCompletedTask(t *testing.T) { runID := "run-ts-3" checkBlockingResult(t, diff --git a/apps/edge/internal/openai/run_result.go b/apps/edge/internal/openai/run_result.go index 19fee31..71916cd 100644 --- a/apps/edge/internal/openai/run_result.go +++ b/apps/edge/internal/openai/run_result.go @@ -10,6 +10,9 @@ import ( ) func collectRunResult(ctx context.Context, stream edgeservice.RunStream, timeout time.Duration) (string, string, *openAIUsage, error) { + if stream.Events == nil { + return "", "", nil, fmt.Errorf("run stream unavailable") + } var contentBuilder strings.Builder var reasoningBuilder strings.Builder var usage *openAIUsage @@ -21,11 +24,18 @@ func collectRunResult(ctx context.Context, stream edgeservice.RunStream, timeout return "", "", nil, ctx.Err() case <-timer.C: return "", "", nil, fmt.Errorf("run timed out") - case nodeEvent := <-stream.NodeEvents: + case nodeEvent, ok := <-stream.NodeEvents: + if !ok { + stream.NodeEvents = nil + continue + } if edgeservice.IsNodeDisconnected(nodeEvent) { return "", "", nil, fmt.Errorf("node disconnected") } - case event := <-stream.Events: + case event, ok := <-stream.Events: + if !ok { + return "", "", nil, fmt.Errorf("run stream closed") + } if event == nil { continue } diff --git a/apps/edge/internal/openai/server_test.go b/apps/edge/internal/openai/server_test.go index d7480b5..3b974d0 100644 --- a/apps/edge/internal/openai/server_test.go +++ b/apps/edge/internal/openai/server_test.go @@ -172,6 +172,29 @@ func TestChatCompletionsStreamsSSE(t *testing.T) { } } +func TestChatCompletionsStreamingReportsClosedRunStream(t *testing.T) { + fake := &fakeRunService{events: make(chan *iop.RunEvent)} + close(fake.events) + + srv := NewServer(config.EdgeOpenAIConf{Adapter: "ollama"}, fake, nil) + req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{ + "model":"stream-model", + "stream":true, + "messages":[{"role":"user","content":"hi"}] + }`)) + w := httptest.NewRecorder() + + srv.handleChatCompletions(w, req) + + if w.Code != http.StatusOK { + t.Fatalf("status: got %d body=%s", w.Code, w.Body.String()) + } + body := w.Body.String() + if !strings.Contains(body, `"message":"run stream closed"`) || !strings.Contains(body, "data: [DONE]") { + t.Fatalf("unexpected SSE error body:\n%s", body) + } +} + func TestChatCompletionsReturnsReasoningContentSeparately(t *testing.T) { fake := &fakeRunService{events: make(chan *iop.RunEvent, 5)} fake.events <- &iop.RunEvent{Type: "reasoning_delta", Delta: "thinking "} @@ -762,3 +785,23 @@ func TestCollectRunResultTimesOut(t *testing.T) { t.Fatal("expected timeout error") } } + +func TestCollectRunResultFailsWhenEventStreamCloses(t *testing.T) { + events := make(chan *iop.RunEvent) + close(events) + handle := &edgeservice.RunHandle{ + RunDispatch: edgeservice.RunDispatch{TimeoutSec: 60}, + RunStream: edgeservice.RunStream{ + Events: events, + NodeEvents: make(chan *iop.EdgeNodeEvent), + }, + } + + _, _, _, err := collectRunResult(context.Background(), handle.Stream(), handle.WaitTimeout()) + if err == nil { + t.Fatal("expected closed stream error") + } + if !strings.Contains(err.Error(), "run stream closed") { + t.Fatalf("expected run stream closed error, got %v", err) + } +} diff --git a/apps/edge/internal/openai/stream.go b/apps/edge/internal/openai/stream.go index d57465a..7e758ff 100644 --- a/apps/edge/internal/openai/stream.go +++ b/apps/edge/internal/openai/stream.go @@ -56,16 +56,28 @@ func (s *Server) streamChatCompletion(w http.ResponseWriter, r *http.Request, re }() stream := handle.Stream() + if stream.Events == nil { + writeSSEError(w, flusher, "run stream unavailable") + return + } for { select { case <-r.Context().Done(): return - case nodeEvent := <-stream.NodeEvents: + case nodeEvent, ok := <-stream.NodeEvents: + if !ok { + stream.NodeEvents = nil + continue + } if edgeservice.IsNodeDisconnected(nodeEvent) { writeSSEError(w, flusher, "node disconnected") return } - case event := <-stream.Events: + case event, ok := <-stream.Events: + if !ok { + writeSSEError(w, flusher, "run stream closed") + return + } if event == nil { continue } diff --git a/apps/edge/internal/service/run_dispatch.go b/apps/edge/internal/service/run_dispatch.go index 295628e..96e6ae4 100644 --- a/apps/edge/internal/service/run_dispatch.go +++ b/apps/edge/internal/service/run_dispatch.go @@ -3,6 +3,7 @@ package service import ( "context" "fmt" + "sync" "time" "google.golang.org/protobuf/types/known/structpb" @@ -59,12 +60,13 @@ type RunResult interface { type RunHandle struct { RunDispatch RunStream - close func() + closeOnce sync.Once + close func() } func (h *RunHandle) Close() { if h != nil && h.close != nil { - h.close() + h.closeOnce.Do(h.close) } } diff --git a/apps/edge/internal/service/run_dispatch_internal_test.go b/apps/edge/internal/service/run_dispatch_internal_test.go new file mode 100644 index 0000000..d915f06 --- /dev/null +++ b/apps/edge/internal/service/run_dispatch_internal_test.go @@ -0,0 +1,19 @@ +package service + +import "testing" + +func TestRunHandleCloseIsIdempotent(t *testing.T) { + calls := 0 + handle := &RunHandle{ + close: func() { + calls++ + }, + } + + handle.Close() + handle.Close() + + if calls != 1 { + t.Fatalf("close called %d times, want 1", calls) + } +}