iop/apps/edge/internal/openai/chat_stream_session_test.go

package openai

import (
	"context"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"

	"iop/packages/go/config"

	edgeservice "iop/apps/edge/internal/service"
	eventpkg "iop/packages/go/events"
	iop "iop/proto/gen/iop"
)

// countTerminals reports how many SSE terminal markers a session wrote. Every
// terminal path — finish, fail, and the timeout notice — ends with [DONE], so a
// value other than 1 means the session emitted no terminal or emitted twice.
func countTerminals(body string) int {
	return strings.Count(body, "data: [DONE]")
}

// runStreamSession drives a live SSE session over the supplied events and
// returns the response body.
func runStreamSession(t *testing.T, req chatCompletionRequest, outputPolicy strictOutputPolicy, waitTimeout time.Duration, ctx context.Context, emit func(events chan *iop.RunEvent, nodeEvents chan *iop.EdgeNodeEvent)) string {
	t.Helper()
	srv := NewServer(config.EdgeOpenAIConf{Adapter: "ollama"}, &fakeRunService{}, nil)

	events := make(chan *iop.RunEvent, 8)
	nodeEvents := make(chan *iop.EdgeNodeEvent, 4)
	handle := &fakeRunResultWithTimeout{
		dispatch: edgeservice.RunDispatch{
			RunID:   "run-session",
			NodeID:  "node-1",
			Adapter: "ollama",
			Target:  "llama3",
		},
		stream:      edgeservice.RunStream{Events: events, NodeEvents: nodeEvents},
		waitTimeout: waitTimeout,
	}

	emit(events, nodeEvents)

	httpReq := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil).WithContext(ctx)
	w := httptest.NewRecorder()
	dc := newTestChatDispatchContext(httpReq, req)
	dc.outputPolicy = outputPolicy
	srv.streamChatCompletion(w, dc, handle)
	return w.Body.String()
}

// streamSessionCase is one terminal outcome of the live SSE session.
type streamSessionCase struct {
	name         string
	req          chatCompletionRequest
	outputPolicy strictOutputPolicy
	waitTimeout  time.Duration
	emit         func(events chan *iop.RunEvent, nodeEvents chan *iop.EdgeNodeEvent)
	wantBody     []string
	wantMissing  []string
}

// emitReasoningThenAnswer emits a reasoning delta, a content delta, and a
// complete event — the shape the reasoning visibility rules are decided on.
func emitReasoningThenAnswer(events chan *iop.RunEvent, _ chan *iop.EdgeNodeEvent) {
	events <- &iop.RunEvent{Type: "reasoning_delta", Delta: "thinking hard"}
	events <- &iop.RunEvent{Type: "delta", Delta: "answer"}
	events <- &iop.RunEvent{Type: "complete"}
}

// streamSessionTerminalCases covers how a session ends.
func streamSessionTerminalCases() []streamSessionCase {
	return []streamSessionCase{
		{
			name:        "complete emits content then finish",
			req:         chatCompletionRequest{Model: "llama3"},
			waitTimeout: time.Second,
			emit: func(events chan *iop.RunEvent, _ chan *iop.EdgeNodeEvent) {
				events <- &iop.RunEvent{Type: "delta", Delta: "hello"}
				events <- &iop.RunEvent{Type: "complete", Metadata: map[string]string{"finish_reason": "stop"}}
			},
			wantBody: []string{`"role":"assistant"`, `"content":"hello"`, `"finish_reason":"stop"`},
		},
		{
			name:        "error event surfaces the run error",
			req:         chatCompletionRequest{Model: "llama3"},
			waitTimeout: time.Second,
			emit: func(events chan *iop.RunEvent, _ chan *iop.EdgeNodeEvent) {
				events <- &iop.RunEvent{Type: "error", Error: "boom"}
			},
			wantBody:    []string{"boom", "run_error"},
			wantMissing: []string{`"finish_reason":"stop"`},
		},
		{
			name:        "cancelled event surfaces its message",
			req:         chatCompletionRequest{Model: "llama3"},
			waitTimeout: time.Second,
			emit: func(events chan *iop.RunEvent, _ chan *iop.EdgeNodeEvent) {
				events <- &iop.RunEvent{Type: "cancelled", Message: "stopped by user"}
			},
			wantBody: []string{"stopped by user"},
		},
		{
			name:        "closed stream terminates once",
			req:         chatCompletionRequest{Model: "llama3"},
			waitTimeout: time.Second,
			emit: func(events chan *iop.RunEvent, _ chan *iop.EdgeNodeEvent) {
				close(events)
			},
			wantBody: []string{"run stream closed"},
		},
		{
			name:        "node disconnect terminates once",
			req:         chatCompletionRequest{Model: "llama3"},
			waitTimeout: time.Second,
			emit: func(_ chan *iop.RunEvent, nodeEvents chan *iop.EdgeNodeEvent) {
				nodeEvents <- &iop.EdgeNodeEvent{Type: eventpkg.TypeNodeDisconnected, NodeId: "node-1"}
			},
			wantBody: []string{"node disconnected"},
		},
		{
			name:        "timeout terminates once",
			req:         chatCompletionRequest{Model: "llama3"},
			waitTimeout: 30 * time.Millisecond,
			emit:        func(_ chan *iop.RunEvent, _ chan *iop.EdgeNodeEvent) {},
			wantBody:    []string{"run timed out"},
		},
	}
}

// streamSessionReasoningCases covers what a session exposes of the model's
// reasoning.
func streamSessionReasoningCases() []streamSessionCase {
	strictTrue := true
	return []streamSessionCase{
		{
			name:        "reasoning is exposed by default",
			req:         chatCompletionRequest{Model: "llama3"},
			waitTimeout: time.Second,
			emit:        emitReasoningThenAnswer,
			wantBody:    []string{`"reasoning_content":"thinking hard"`, `"content":"answer"`},
		},
		{
			name:         "strict output hides reasoning the caller did not ask for",
			req:          chatCompletionRequest{Model: "llama3"},
			outputPolicy: strictOutputPolicy{Strict: true},
			waitTimeout:  time.Second,
			emit:         emitReasoningThenAnswer,
			wantBody:     []string{`"content":"answer"`},
			wantMissing:  []string{"thinking hard"},
		},
		{
			name:         "strict output exposes reasoning the caller asked for",
			req:          chatCompletionRequest{Model: "llama3", IncludeReasoning: &strictTrue},
			outputPolicy: strictOutputPolicy{Strict: true},
			waitTimeout:  time.Second,
			emit:         emitReasoningThenAnswer,
			wantBody:     []string{`"reasoning_content":"thinking hard"`},
		},
	}
}

// TestChatStreamSessionTerminalOutcomes pins the terminal contract of the live
// SSE session: each outcome emits exactly one terminal and the accumulated
// content it was supposed to expose.
func TestChatStreamSessionTerminalOutcomes(t *testing.T) {
	cases := append(streamSessionTerminalCases(), streamSessionReasoningCases()...)
	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			body := runStreamSession(t, tc.req, tc.outputPolicy, tc.waitTimeout, context.Background(), tc.emit)

			if got := countTerminals(body); got != 1 {
				t.Fatalf("expected exactly 1 terminal, got %d: %s", got, body)
			}
			for _, want := range tc.wantBody {
				if !strings.Contains(body, want) {
					t.Fatalf("expected body to contain %q, got %s", want, body)
				}
			}
			for _, missing := range tc.wantMissing {
				if strings.Contains(body, missing) {
					t.Fatalf("expected body to omit %q, got %s", missing, body)
				}
			}
		})
	}
}

// TestChatStreamSessionCallerCancelWritesNoTerminal pins that a caller who
// disconnects gets no SSE body: the run is cancelled node-side instead.
func TestChatStreamSessionCallerCancelWritesNoTerminal(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	cancel()

	body := runStreamSession(t, chatCompletionRequest{Model: "llama3"}, strictOutputPolicy{}, time.Second, ctx,
		func(_ chan *iop.RunEvent, _ chan *iop.EdgeNodeEvent) {})

	if got := countTerminals(body); got != 0 {
		t.Fatalf("expected no terminal for a cancelled caller, got %d: %s", got, body)
	}
}