package admin import ( "bytes" "encoding/json" "io" "net/http" "net/http/httptest" "os" "path/filepath" "strings" "testing" "ds2api/internal/devcapture" ) type stubOpenAIChatCaller struct{} func (stubOpenAIChatCaller) ChatCompletions(w http.ResponseWriter, _ *http.Request) { store := devcapture.Global() session := store.Start("deepseek_completion", "https://chat.deepseek.com/api/v0/chat/completion", "acct-test", map[string]any{"model": "deepseek-chat"}) raw := io.NopCloser(strings.NewReader( "data: {\"v\":\"hello [reference:1]\"}\n\n" + "data: {\"v\":\"FINISHED\",\"p\":\"response/status\"}\n\n", )) if session != nil { raw = session.WrapBody(raw, http.StatusOK) } _, _ = io.ReadAll(raw) _ = raw.Close() w.Header().Set("Content-Type", "text/event-stream") w.WriteHeader(http.StatusOK) _, _ = io.WriteString(w, "data: {\"choices\":[{\"delta\":{\"content\":\"hello\"},\"index\":0}],\"created\":1,\"id\":\"id\",\"model\":\"m\",\"object\":\"chat.completion.chunk\"}\n\n") } type stubOpenAIChatCallerWithContinuations struct{} func (stubOpenAIChatCallerWithContinuations) ChatCompletions(w http.ResponseWriter, _ *http.Request) { recordCapturedResponse("deepseek_completion", "https://chat.deepseek.com/api/v0/chat/completion", http.StatusOK, map[string]any{"model": "deepseek-chat"}, "data: {\"v\":\"hello [reference:1]\"}\n\n"+"data: [DONE]\n\n") recordCapturedResponse("deepseek_continue", "https://chat.deepseek.com/api/v0/chat/continue", http.StatusOK, map[string]any{"chat_session_id": "session-1", "message_id": 2}, "data: {\"v\":\"continued\"}\n\n"+"data: [DONE]\n\n") w.Header().Set("Content-Type", "text/event-stream") w.WriteHeader(http.StatusOK) _, _ = io.WriteString(w, "data: {\"choices\":[{\"delta\":{\"content\":\"hello continued\"},\"index\":0}],\"created\":1,\"id\":\"id\",\"model\":\"m\",\"object\":\"chat.completion.chunk\"}\n\n") } type stubOpenAIChatCallerWithoutCapture struct{} func (stubOpenAIChatCallerWithoutCapture) ChatCompletions(w http.ResponseWriter, _ *http.Request) { w.Header().Set("Content-Type", "text/event-stream") w.WriteHeader(http.StatusOK) _, _ = io.WriteString(w, "data: {\"choices\":[{\"delta\":{\"content\":\"hello\"},\"index\":0}],\"created\":1,\"id\":\"id\",\"model\":\"m\",\"object\":\"chat.completion.chunk\"}\n\n") } func recordCapturedResponse(label, rawURL string, statusCode int, request any, body string) { store := devcapture.Global() session := store.Start(label, rawURL, "acct-test", request) raw := io.NopCloser(strings.NewReader(body)) if session != nil { raw = session.WrapBody(raw, statusCode) } _, _ = io.ReadAll(raw) _ = raw.Close() } func TestCaptureRawSampleWritesPersistentSample(t *testing.T) { t.Setenv("DS2API_RAW_STREAM_SAMPLE_ROOT", t.TempDir()) devcapture.Global().Clear() defer devcapture.Global().Clear() h := &Handler{OpenAI: stubOpenAIChatCaller{}} reqBody := `{ "sample_id":"My Sample 01", "api_key":"local-key", "model":"deepseek-chat", "message":"广州天气", "stream":true }` rec := httptest.NewRecorder() req := httptest.NewRequest(http.MethodPost, "/admin/dev/raw-samples/capture", strings.NewReader(reqBody)) h.captureRawSample(rec, req) if rec.Code != http.StatusOK { t.Fatalf("expected 200, got %d body=%s", rec.Code, rec.Body.String()) } if got := rec.Header().Get("X-Ds2-Sample-Id"); got != "my-sample-01" { t.Fatalf("expected sample id header my-sample-01, got %q", got) } if got := rec.Header().Get("X-Ds2-Sample-Upstream"); got != filepath.Join(os.Getenv("DS2API_RAW_STREAM_SAMPLE_ROOT"), "my-sample-01", "upstream.stream.sse") { t.Fatalf("unexpected sample upstream header: %q", got) } if !strings.Contains(rec.Body.String(), `"content":"hello"`) { t.Fatalf("expected proxied openai output, got %s", rec.Body.String()) } sampleDir := filepath.Join(os.Getenv("DS2API_RAW_STREAM_SAMPLE_ROOT"), "my-sample-01") if _, err := os.Stat(sampleDir); err != nil { t.Fatalf("sample dir missing: %v", err) } metaBytes, err := os.ReadFile(filepath.Join(sampleDir, "meta.json")) if err != nil { t.Fatalf("read meta: %v", err) } var meta map[string]any if err := json.Unmarshal(metaBytes, &meta); err != nil { t.Fatalf("decode meta: %v", err) } if meta["sample_id"] != "my-sample-01" { t.Fatalf("unexpected meta sample_id: %#v", meta["sample_id"]) } capture, _ := meta["capture"].(map[string]any) if capture == nil { t.Fatalf("missing capture meta: %#v", meta) } if got := int(capture["response_bytes"].(float64)); got == 0 { t.Fatalf("expected capture bytes to be recorded, got %#v", capture) } if _, ok := meta["processed"]; ok { t.Fatalf("unexpected processed meta: %#v", meta["processed"]) } } func TestCaptureRawSampleCombinesContinuationCaptures(t *testing.T) { t.Setenv("DS2API_RAW_STREAM_SAMPLE_ROOT", t.TempDir()) devcapture.Global().Clear() defer devcapture.Global().Clear() h := &Handler{OpenAI: stubOpenAIChatCallerWithContinuations{}} reqBody := `{ "sample_id":"My Sample 02", "api_key":"local-key", "model":"deepseek-chat", "message":"广州天气", "stream":true }` rec := httptest.NewRecorder() req := httptest.NewRequest(http.MethodPost, "/admin/dev/raw-samples/capture", strings.NewReader(reqBody)) h.captureRawSample(rec, req) if rec.Code != http.StatusOK { t.Fatalf("expected 200, got %d body=%s", rec.Code, rec.Body.String()) } sampleDir := filepath.Join(os.Getenv("DS2API_RAW_STREAM_SAMPLE_ROOT"), "my-sample-02") upstreamBytes, err := os.ReadFile(filepath.Join(sampleDir, "upstream.stream.sse")) if err != nil { t.Fatalf("read upstream: %v", err) } upstream := string(upstreamBytes) if !strings.Contains(upstream, "hello [reference:1]") { t.Fatalf("expected initial capture in combined upstream, got %s", upstream) } if !strings.Contains(upstream, "continued") { t.Fatalf("expected continuation capture in combined upstream, got %s", upstream) } if strings.Index(upstream, "hello [reference:1]") > strings.Index(upstream, "continued") { t.Fatalf("expected initial capture before continuation, got %s", upstream) } metaBytes, err := os.ReadFile(filepath.Join(sampleDir, "meta.json")) if err != nil { t.Fatalf("read meta: %v", err) } var meta map[string]any if err := json.Unmarshal(metaBytes, &meta); err != nil { t.Fatalf("decode meta: %v", err) } capture, _ := meta["capture"].(map[string]any) if capture == nil { t.Fatalf("missing capture meta: %#v", meta) } if got := int(capture["response_bytes"].(float64)); got != len(upstreamBytes) { t.Fatalf("expected combined response_bytes %d, got %#v", len(upstreamBytes), capture["response_bytes"]) } rounds, _ := capture["rounds"].([]any) if len(rounds) != 2 { t.Fatalf("expected 2 capture rounds, got %d: %#v", len(rounds), capture) } r0, _ := rounds[0].(map[string]any) r1, _ := rounds[1].(map[string]any) if r0["label"] != "deepseek_completion" { t.Fatalf("expected first round label deepseek_completion, got %v", r0["label"]) } if r1["label"] != "deepseek_continue" { t.Fatalf("expected second round label deepseek_continue, got %v", r1["label"]) } } func TestCaptureRawSampleReturnsErrorWhenNoNewCaptureRecorded(t *testing.T) { root := t.TempDir() t.Setenv("DS2API_RAW_STREAM_SAMPLE_ROOT", root) devcapture.Global().Clear() defer devcapture.Global().Clear() recordCapturedResponse("preexisting", "https://chat.deepseek.com/api/v0/chat/completion", http.StatusOK, map[string]any{"model": "deepseek-chat"}, "data: {\"v\":\"old\"}\n\n") h := &Handler{OpenAI: stubOpenAIChatCallerWithoutCapture{}} reqBody := `{ "sample_id":"My Sample 03", "api_key":"local-key", "model":"deepseek-chat", "message":"广州天气", "stream":true }` rec := httptest.NewRecorder() req := httptest.NewRequest(http.MethodPost, "/admin/dev/raw-samples/capture", strings.NewReader(reqBody)) h.captureRawSample(rec, req) if rec.Code != http.StatusInternalServerError { t.Fatalf("expected 500, got %d body=%s", rec.Code, rec.Body.String()) } if !strings.Contains(rec.Body.String(), "no upstream capture was recorded") { t.Fatalf("expected no-capture error, got %s", rec.Body.String()) } if _, err := os.Stat(filepath.Join(root, "my-sample-03")); !os.IsNotExist(err) { t.Fatalf("expected no sample dir to be created, stat err=%v", err) } } func TestCombineCaptureBodiesPreservesOrderAndSeparators(t *testing.T) { entries := []devcapture.Entry{ {ResponseBody: "first"}, {ResponseBody: "second"}, } got := combineCaptureBodies(entries) if !bytes.Equal(got, []byte("first\nsecond")) { t.Fatalf("unexpected combined body: %q", string(got)) } } func TestQueryRawSampleCapturesGroupsBySessionAndMatchesQuestion(t *testing.T) { devcapture.Global().Clear() defer devcapture.Global().Clear() recordCapturedResponse( "deepseek_completion", "https://chat.deepseek.com/api/v0/chat/completion", http.StatusOK, map[string]any{ "chat_session_id": "session-query-1", "prompt": "用户问题:广州天气怎么样?", }, "data: {\"v\":\"先看天气\"}\n\n", ) recordCapturedResponse( "deepseek_continue", "https://chat.deepseek.com/api/v0/chat/continue", http.StatusOK, map[string]any{ "chat_session_id": "session-query-1", "message_id": 2, }, "data: {\"v\":\"再补充一点\"}\n\n", ) h := &Handler{} rec := httptest.NewRecorder() req := httptest.NewRequest(http.MethodGet, "/admin/dev/raw-samples/query?q=广州天气", nil) h.queryRawSampleCaptures(rec, req) if rec.Code != http.StatusOK { t.Fatalf("expected 200, got %d body=%s", rec.Code, rec.Body.String()) } var out map[string]any if err := json.Unmarshal(rec.Body.Bytes(), &out); err != nil { t.Fatalf("decode failed: %v", err) } items, _ := out["items"].([]any) if len(items) != 1 { t.Fatalf("expected 1 item, got %d body=%s", len(items), rec.Body.String()) } item, _ := items[0].(map[string]any) if item["chain_key"] != "session:session-query-1" { t.Fatalf("unexpected chain key: %#v", item["chain_key"]) } if int(item["round_count"].(float64)) != 2 { t.Fatalf("expected 2 rounds, got %#v", item["round_count"]) } reqPreview, _ := item["request_preview"].(string) if !strings.Contains(reqPreview, "广州天气") { t.Fatalf("expected request preview to contain query, got %q", reqPreview) } } func TestSaveRawSampleFromCapturesPersistsSelectedChain(t *testing.T) { root := t.TempDir() t.Setenv("DS2API_RAW_STREAM_SAMPLE_ROOT", root) devcapture.Global().Clear() defer devcapture.Global().Clear() recordCapturedResponse( "deepseek_completion", "https://chat.deepseek.com/api/v0/chat/completion", http.StatusOK, map[string]any{ "chat_session_id": "session-save-1", "prompt": "请回答深圳天气", }, "data: {\"v\":\"第一段\"}\n\n", ) recordCapturedResponse( "deepseek_continue", "https://chat.deepseek.com/api/v0/chat/continue", http.StatusOK, map[string]any{ "chat_session_id": "session-save-1", "message_id": 2, }, "data: {\"v\":\"第二段\"}\n\n", ) h := &Handler{} rec := httptest.NewRecorder() reqBody := `{"query":"深圳天气","sample_id":"saved-from-memory"}` req := httptest.NewRequest(http.MethodPost, "/admin/dev/raw-samples/save", strings.NewReader(reqBody)) h.saveRawSampleFromCaptures(rec, req) if rec.Code != http.StatusOK { t.Fatalf("expected 200, got %d body=%s", rec.Code, rec.Body.String()) } var out map[string]any if err := json.Unmarshal(rec.Body.Bytes(), &out); err != nil { t.Fatalf("decode failed: %v", err) } if out["sample_id"] != "saved-from-memory" { t.Fatalf("unexpected sample id: %#v", out["sample_id"]) } if int(out["round_count"].(float64)) != 2 { t.Fatalf("expected round_count=2, got %#v", out["round_count"]) } sampleDir := filepath.Join(root, "saved-from-memory") upstreamBytes, err := os.ReadFile(filepath.Join(sampleDir, "upstream.stream.sse")) if err != nil { t.Fatalf("read upstream: %v", err) } upstream := string(upstreamBytes) if !strings.Contains(upstream, "第一段") || !strings.Contains(upstream, "第二段") { t.Fatalf("expected combined upstream, got %q", upstream) } metaBytes, err := os.ReadFile(filepath.Join(sampleDir, "meta.json")) if err != nil { t.Fatalf("read meta: %v", err) } var meta map[string]any if err := json.Unmarshal(metaBytes, &meta); err != nil { t.Fatalf("decode meta: %v", err) } reqMeta, _ := meta["request"].(map[string]any) if fieldString(reqMeta, "chat_session_id") != "session-save-1" { t.Fatalf("expected request to come from selected chain, got %#v", meta["request"]) } }