go-sip/internal/ai/protocol_turn_test.go

package ai

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync/atomic"
	"testing"
	"time"

	"github.com/coder/websocket"
)

func TestProtocolModelsCompleteTurnWithoutPerModelCode(t *testing.T) {
	testProtocolResampledTurn(t, 16000, 16000)
}

func TestTaskASRTTSSampleRateCombinations(t *testing.T) {
	for _, asrRate := range []int{8000, 16000, 24000, 48000} {
		for _, ttsRate := range []int{8000, 16000, 24000, 48000} {
			t.Run(fmt.Sprintf("full_ai_asr%d_tts%d", asrRate, ttsRate), func(t *testing.T) {
				testProtocolResampledTurn(t, asrRate, ttsRate)
			})
		}
	}
}

func testProtocolResampledTurn(t *testing.T, asrRate, ttsRate int) {
	t.Helper()
	task, providers := protocolTask(t, TTSProtocolDashScopeTask)
	task = changeCurrentAgent(t, task, func(a map[string]any) {
		a["conversation"].(map[string]any)["opening"] = "Hello."
		a["conversation"].(map[string]any)["hangup_keywords"] = []any{}
		a["asr"].(map[string]any)["params"].(map[string]any)["sample_rate"] = asrRate
		a["tts"].(map[string]any)["params"].(map[string]any)["sample_rate"] = ttsRate
	})
	bound, err := Bind(task, providers)
	if err != nil {
		t.Fatal(err)
	}
	var asrCalls, ttsCalls, llmCalls atomic.Int32
	// A 100-ms constant PCM16 signal independently verifies sample values and
	// duration, rather than calculating expectations with the production resampler.
	// 1008 is exactly representable in PCMA, so both rates and values can be
	// checked independently through the real RTP codec boundary.
	expected := bytes.Repeat([]byte{0xf0, 0x03}, 1600)
	ttsAudio := bytes.Repeat([]byte{0xf0, 0x03}, ttsRate/10)
	speech := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		c, err := websocket.Accept(w, r, nil)
		if err != nil {
			t.Error(err)
			return
		}
		defer c.CloseNow()
		ctx, cancel := context.WithTimeout(r.Context(), 2*time.Second)
		defer cancel()
		cmd := mockRead(t, c, ctx)
		mockEvent(c, ctx, cmd.Header.TaskID, "task-started", nil)
		switch cmd.Payload["task"] {
		case "asr":
			asrCalls.Add(1)
			if cmd.Payload["model"] != "unlisted-asr-model-2029" {
				t.Error("ASR model replaced")
			}
			if got := cmd.Payload["parameters"].(map[string]any)["sample_rate"]; got != float64(asrRate) {
				t.Errorf("ASR declared rate = %v, want %d", got, asrRate)
			}
			var received []byte
			for {
				kind, chunk, err := c.Read(ctx)
				if err != nil {
					t.Error(err)
					return
				}
				if kind == websocket.MessageText {
					break
				}
				received = append(received, chunk...)
			}
			if want := bytes.Repeat([]byte{0xf0, 0x03}, asrRate/10); !bytes.Equal(received, want) {
				t.Errorf("ASR actual PCM: %d bytes, want %d with unchanged sample values", len(received), len(want))
			}
			mockEvent(c, ctx, cmd.Header.TaskID, "result-generated", map[string]any{"output": map[string]any{"sentence": map[string]any{"text": "Tell me more.", "sentence_end": true, "begin_time": 0, "end_time": 100}}})
		case "tts":
			ttsCalls.Add(1)
			if cmd.Payload["model"] != "unlisted-tts-model-2029" || cmd.Payload["parameters"].(map[string]any)["voice"] != "unlisted-voice" {
				t.Error("TTS model/voice replaced")
			}
			if got := cmd.Payload["parameters"].(map[string]any)["sample_rate"]; got != float64(ttsRate) {
				t.Errorf("TTS declared rate = %v, want %d", got, ttsRate)
			}
			mockRead(t, c, ctx)
			finish := mockRead(t, c, ctx)
			if finish.Header.Action != "finish-task" {
				t.Error("TTS lifecycle changed")
			}
			c.Write(ctx, websocket.MessageBinary, ttsAudio)
		default:
			t.Error("unknown role")
			return
		}
		mockEvent(c, ctx, cmd.Header.TaskID, "task-finished", nil)
	}))
	defer speech.Close()
	llm := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		llmCalls.Add(1)
		var request map[string]any
		if r.Method != "POST" || r.URL.Path != "/v1/chat/completions" || json.NewDecoder(r.Body).Decode(&request) != nil || request["model"] != "unlisted-chat-model-2029" {
			t.Error("LLM model/protocol changed")
		}
		w.Header().Set("Content-Type", "application/json")
		fmt.Fprint(w, `{"id":"mock","object":"chat.completion","created":0,"model":"unlisted-chat-model-2029","choices":[{"index":0,"message":{"role":"assistant","content":"Generic response."},"finish_reason":"stop"}]}`)
	}))
	defer llm.Close()
	bound.ASR.Provider.Endpoint = "ws" + strings.TrimPrefix(speech.URL, "http")
	bound.TTS.Provider.WSEndpoint = bound.ASR.Provider.Endpoint
	bound.LLM.Provider.Endpoint = llm.URL + "/v1"
	call, err := NewCall(bound, func(context.Context) error { t.Error("ordinary turn attempted hangup"); return nil })
	if err != nil {
		t.Fatal(err)
	}
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	opening, err := call.Open(ctx)
	if err != nil || !bytes.Equal(opening, expected) {
		t.Fatalf("opening: %v", err)
	}
	if played := protocolPCMAViaRTP(t, ctx, opening); !bytes.Equal(played, expected) {
		t.Fatal("opening changed across Asterisk PCMA media boundary")
	}
	userPCM := protocolPCMAViaRTP(t, ctx, expected)
	result, err := call.RunTurn(ctx, userPCM)
	if err != nil || result.Transcript != "Tell me more." || result.Reply != "Generic response." || !bytes.Equal(result.AudioPCM16, expected) || result.EndedByKeyword {
		t.Fatalf("generic model turn: %+v %v", result, err)
	}
	if played := protocolPCMAViaRTP(t, ctx, result.AudioPCM16); !bytes.Equal(played, expected) {
		t.Fatal("AI reply changed across Asterisk PCMA media boundary")
	}
	if asrCalls.Load() != 1 || llmCalls.Load() != 1 || ttsCalls.Load() != 2 {
		t.Fatalf("unexpected/retried requests: ASR=%d LLM=%d TTS=%d", asrCalls.Load(), llmCalls.Load(), ttsCalls.Load())
	}
}