fix: music poll resilience + hostile-URL guard + embed dup-index (gadfly round 1)

- pollResult tolerates up to 5 CONSECUTIVE bad polls (transport blip, unparseable payload, task momentarily absent) instead of killing a multi-minute exclusive-GPU job on the first hiccup; only status=2, a failure run, or ctx deadline aborts - server-supplied result.File must be server-relative; combined with the upstreamPath dot-dot/scheme rejection this stops a hostile upstream from steering the follow-up GET at other proxy endpoints (test: ../../api/models/unload refused) - WithSteps(<=0) rejected; embed responses repeating an index rejected; musicFormatMIME now wraps speechMIME (one format table, wav32 normalized); poll interval is a test-shrinkable var (CI no longer burns 2s+ per music test); parens + comments per review Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-07-13 00:52:06 -04:00
parent cacce61ddc
commit cf2d83f157
3 changed files with 143 additions and 35 deletions
@@ -75,6 +75,10 @@ func (m *embedModel) Embed(ctx context.Context, req embeddings.EmbedRequest, opt
 			return nil, &llm.APIError{Provider: m.p.name, Model: m.id,
 				Message: fmt.Sprintf("embeddings response entry index %d invalid or empty", d.Index)}
 		}
+		if vectors[d.Index] != nil {
+			return nil, &llm.APIError{Provider: m.p.name, Model: m.id,
+				Message: fmt.Sprintf("embeddings response repeats index %d", d.Index)}
+		}
 		vectors[d.Index] = d.Embedding
 	}
 	for i, v := range vectors {
@@ -25,8 +25,14 @@ import (

 // musicPollInterval is the delay between query_result polls. Long enough to
 // be polite to the queue, short enough that a ~10s xl-turbo song isn't
-// dominated by poll latency.
-const musicPollInterval = 2 * time.Second
+// dominated by poll latency. A var so tests can shrink it.
+var musicPollInterval = 2 * time.Second
+
+// musicPollMaxConsecutiveFailures bounds how many consecutive BAD polls
+// (transport error, unparseable payload, task momentarily absent) are
+// tolerated before aborting. A multi-minute GPU job must not die to one
+// blip; a genuinely broken upstream still fails within ~5 intervals.
+const musicPollMaxConsecutiveFailures = 5

 // MusicModel implements musicgen.Provider. The id selects which upstream
 // llama-swap loads.
@@ -44,7 +50,8 @@ type musicModel struct {
 }

 // releaseTaskRequest is the ACE-Step POST /release_task shape. audio_duration
-// is the v1 param name; verify against docs/en/API.md at smoke time — an
+// is the v1 param name; verify against the upstream ACE-Step-1.5 repo's
+// docs/en/API.md at smoke time — an
 // unknown field is ignored upstream, degrading to the default clip length,
 // never an error.
 type releaseTaskRequest struct {
@@ -65,20 +72,16 @@ type queryItem struct {
 	Result string `json:"result"`
 }

-// musicFormatMIME maps ACE-Step's audio_format values to MIME types.
-func musicFormatMIME(format string) string {
-	switch strings.ToLower(strings.TrimSpace(format)) {
-	case "wav", "wav32":
-		return "audio/wav"
-	case "flac":
-		return "audio/flac"
-	case "opus":
-		return "audio/ogg"
-	case "aac":
-		return "audio/aac"
-	default: // "", "mp3"
-		return "audio/mpeg"
+// musicFormatMIME resolves the clip MIME from the response Content-Type
+// or the requested format, reusing speechMIME's table (one format->MIME
+// map for the whole provider). wav32 is ACE-Step-specific: normalize it
+// to wav before the shared lookup.
+func musicFormatMIME(contentType, format string) string {
+	format = strings.ToLower(strings.TrimSpace(format))
+	if format == "wav32" {
+		format = "wav"
 	}
+	return speechMIME(contentType, format)
 }

 // Generate implements musicgen.Model.
@@ -90,6 +93,9 @@ func (m *musicModel) Generate(ctx context.Context, req musicgen.Request, opts ..
 	if req.DurationSeconds < 0 {
 		return nil, fmt.Errorf("%w: duration must be >= 0, got %d", llm.ErrUnsupported, req.DurationSeconds)
 	}
+	if req.Steps != nil && *req.Steps <= 0 {
+		return nil, fmt.Errorf("%w: inference steps must be > 0, got %d", llm.ErrUnsupported, *req.Steps)
+	}

 	taskID, err := m.releaseTask(ctx, req)
 	if err != nil {
@@ -139,7 +145,12 @@ func (m *musicModel) releaseTask(ctx context.Context, req musicgen.Request) (str
 }

 // pollResult polls query_result until the task succeeds, fails, or ctx
-// expires.
+// expires. Transient trouble — a transport blip, a momentarily
+// unparseable payload, the task briefly absent from the response — is
+// tolerated up to musicPollMaxConsecutiveFailures in a row: a
+// multi-minute exclusive-GPU job must not die to one flaky poll. Only an
+// explicit status=2, a run of consecutive failures, or the ctx deadline
+// aborts.
 func (m *musicModel) pollResult(ctx context.Context, taskID string) (*queryItem, error) {
 	path, err := upstreamPath(m.id, "/query_result")
 	if err != nil {
@@ -148,22 +159,29 @@ func (m *musicModel) pollResult(ctx context.Context, taskID string) (*queryItem,
 	body := map[string]any{"task_id_list": []string{taskID}}
 	ticker := time.NewTicker(musicPollInterval)
 	defer ticker.Stop()
+	failures := 0
+	var lastErr error
 	for {
-		// Tolerant envelope: items under "data" or a bare array.
-		var raw json.RawMessage
-		if err := m.p.doJSON(ctx, http.MethodPost, path, m.id, body, &raw); err != nil {
-			return nil, err
-		}
-		item, err := findQueryItem(raw, taskID)
-		if err != nil {
-			return nil, &llm.APIError{Provider: m.p.name, Model: m.id, Message: err.Error()}
-		}
-		switch item.Status {
-		case 1:
+		item, pollErr := m.pollOnce(ctx, path, body, taskID)
+		switch {
+		case pollErr != nil:
+			// Ctx expiry is never transient — bail with the deadline error.
+			if ctx.Err() != nil {
+				return nil, fmt.Errorf("llama-swap: music generation: %w", ctx.Err())
+			}
+			failures++
+			lastErr = pollErr
+			if failures >= musicPollMaxConsecutiveFailures {
+				return nil, &llm.APIError{Provider: m.p.name, Model: m.id,
+					Message: fmt.Sprintf("music poll failed %d times in a row: %v", failures, lastErr)}
+			}
+		case item.Status == 1:
 			return item, nil
-		case 2:
+		case item.Status == 2:
 			return nil, &llm.APIError{Provider: m.p.name, Model: m.id,
 				Message: "music generation failed upstream: " + truncateForError([]byte(item.Result))}
+		default:
+			failures = 0 // healthy queued/running poll
 		}
 		select {
 		case <-ctx.Done():
@@ -173,6 +191,16 @@ func (m *musicModel) pollResult(ctx context.Context, taskID string) (*queryItem,
 	}
 }

+// pollOnce performs one query_result round trip and locates the task.
+func (m *musicModel) pollOnce(ctx context.Context, path string, body any, taskID string) (*queryItem, error) {
+	// Tolerant envelope: items under "data" or a bare array.
+	var raw json.RawMessage
+	if err := m.p.doJSON(ctx, http.MethodPost, path, m.id, body, &raw); err != nil {
+		return nil, err
+	}
+	return findQueryItem(raw, taskID)
+}
+
 // findQueryItem digs the task's entry out of the query_result payload,
 // tolerating {"data": [...]}, {"data": {...}}, and bare-array envelopes.
 func findQueryItem(raw json.RawMessage, taskID string) (*queryItem, error) {
@@ -192,7 +220,10 @@ func findQueryItem(raw json.RawMessage, taskID string) (*queryItem, error) {
 		items = []queryItem{one}
 	}
 	for i := range items {
-		if items[i].TaskID == taskID || items[i].TaskID == "" && len(items) == 1 {
+		// Single-item responses without a task_id echo are assumed to be
+		// ours — we only ever poll one task; a mismatch surfaces as a
+		// transient miss and is retried by the caller.
+		if items[i].TaskID == taskID || (items[i].TaskID == "" && len(items) == 1) {
 			return &items[i], nil
 		}
 	}
@@ -209,7 +240,12 @@ func (m *musicModel) fetchResult(ctx context.Context, format string, item *query
 			Message: "music result blob missing file URL: " + truncateForError([]byte(item.Result))}
 	}
 	// The file URL is server-relative (e.g. "/v1/audio?path=..."); route it
-	// back through the same upstream.
+	// back through the same upstream. upstreamPath additionally refuses
+	// dot-dot/scheme smuggling in this SERVER-SUPPLIED value.
+	if !strings.HasPrefix(result.File, "/") {
+		return nil, &llm.APIError{Provider: m.p.name, Model: m.id,
+			Message: "music result file URL is not server-relative: " + truncateForError([]byte(result.File))}
+	}
 	path, err := upstreamPath(m.id, result.File)
 	if err != nil {
 		return nil, err
@@ -221,10 +257,7 @@ func (m *musicModel) fetchResult(ctx context.Context, format string, item *query
 	if len(raw) == 0 {
 		return nil, &llm.APIError{Provider: m.p.name, Model: m.id, Message: "music response contained no audio"}
 	}
-	mimeType := mimeFromContentType(contentType, "audio/")
-	if mimeType == "" {
-		mimeType = musicFormatMIME(format)
-	}
+	mimeType := musicFormatMIME(contentType, format)
 	return &musicgen.Result{
 		Audio: musicgen.Audio{Data: raw, MIME: mimeType},
 		Raw:   json.RawMessage(item.Result),
@@ -51,6 +51,10 @@ func aceStepStub(t *testing.T, mp3 []byte) (*httptest.Server, *atomic.Int32) {
 }

 func TestMusicGenerate(t *testing.T) {
+	old := musicPollInterval
+	musicPollInterval = 5 * time.Millisecond
+	defer func() { musicPollInterval = old }()
+
 	mp3 := []byte("ID3fakeaudio")
 	srv, polls := aceStepStub(t, mp3)
 	defer srv.Close()
@@ -79,6 +83,10 @@ func TestMusicGenerate(t *testing.T) {
 }

 func TestMusicGenerateUpstreamFailure(t *testing.T) {
+	old := musicPollInterval
+	musicPollInterval = 5 * time.Millisecond
+	defer func() { musicPollInterval = old }()
+
 	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		switch r.URL.Path {
 		case "/upstream/musicgen-acestep/release_task":
@@ -222,3 +230,66 @@ func TestInstructedQuery(t *testing.T) {
 		t.Errorf("InstructedQuery = %q", got)
 	}
 }
+
+func TestMusicGenerateSurvivesTransientPollFailures(t *testing.T) {
+	old := musicPollInterval
+	musicPollInterval = 5 * time.Millisecond
+	defer func() { musicPollInterval = old }()
+
+	mp3 := []byte("ID3fakeaudio")
+	var polls atomic.Int32
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		switch r.URL.Path {
+		case "/upstream/musicgen-acestep/release_task":
+			_, _ = w.Write([]byte(`{"data": {"task_id": "t-3"}}`))
+		case "/upstream/musicgen-acestep/query_result":
+			switch polls.Add(1) {
+			case 1:
+				w.WriteHeader(http.StatusBadGateway) // transient transport blip
+			case 2:
+				_, _ = w.Write([]byte(`{"data": []}`)) // task momentarily absent
+			default:
+				_, _ = w.Write([]byte(`{"data": [{"task_id": "t-3", "status": 1,
+					"result": "{\"file\": \"/v1/audio?path=out.mp3\"}"}]}`))
+			}
+		case "/upstream/musicgen-acestep/v1/audio":
+			w.Header().Set("Content-Type", "audio/mpeg")
+			_, _ = w.Write(mp3)
+		}
+	}))
+	defer srv.Close()
+
+	p := New(WithBaseURL(srv.URL), WithHTTPClient(srv.Client()))
+	mm, _ := p.MusicModel("musicgen-acestep")
+	res, err := mm.Generate(context.Background(), musicgen.Request{Prompt: "p"})
+	if err != nil {
+		t.Fatalf("Generate should survive 2 transient failures: %v", err)
+	}
+	if len(res.Audio.Data) == 0 {
+		t.Fatal("no audio")
+	}
+}
+
+func TestMusicGenerateRejectsHostileFileURL(t *testing.T) {
+	old := musicPollInterval
+	musicPollInterval = 5 * time.Millisecond
+	defer func() { musicPollInterval = old }()
+
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		switch r.URL.Path {
+		case "/upstream/musicgen-acestep/release_task":
+			_, _ = w.Write([]byte(`{"data": {"task_id": "t-4"}}`))
+		case "/upstream/musicgen-acestep/query_result":
+			_, _ = w.Write([]byte(`{"data": [{"task_id": "t-4", "status": 1,
+				"result": "{\"file\": \"/v1/audio?path=../../api/models/unload\"}"}]}`))
+		}
+	}))
+	defer srv.Close()
+
+	p := New(WithBaseURL(srv.URL), WithHTTPClient(srv.Client()))
+	mm, _ := p.MusicModel("musicgen-acestep")
+	_, err := mm.Generate(context.Background(), musicgen.Request{Prompt: "p"})
+	if err == nil {
+		t.Fatal("dot-dot result file URL accepted")
+	}
+}