majordomo/provider/llamaswap/music.go

// music.go implements musicgen.Provider against an ACE-Step-1.5-style API
// server reached through llama-swap's /upstream passthrough (ADR-0021):
//
//	POST /upstream/<id>/release_task   {prompt, lyrics, ...} -> {task_id}
//	POST /upstream/<id>/query_result   {task_id_list: [...]} -> status+result
//	GET  /upstream/<id>/<result file URL>                    -> audio bytes
//
// The backend is an async job queue; Generate wraps it into the blocking
// one-call contract by polling, so a context deadline is the caller's
// budget for the whole job (mort's tool timeout sits well under its
// agent-runtime ceiling for exactly this reason).
package llamaswap

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"strings"
	"time"

	"gitea.stevedudenhoeffer.com/steve/majordomo/llm"
	"gitea.stevedudenhoeffer.com/steve/majordomo/musicgen"
)

// 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. 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.
func (p *Provider) MusicModel(id string, opts ...musicgen.ModelOption) (musicgen.Model, error) {
	if err := p.requireBaseURL(); err != nil {
		return nil, err
	}
	_ = musicgen.ApplyModelOptions(opts)
	return &musicModel{p: p, id: id}, nil
}

type musicModel struct {
	p  *Provider
	id string
}

// releaseTaskRequest is the ACE-Step POST /release_task shape. audio_duration
// 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 {
	Prompt         string `json:"prompt"`
	Lyrics         string `json:"lyrics,omitempty"`
	AudioFormat    string `json:"audio_format,omitempty"`
	TaskType       string `json:"task_type"`
	AudioDuration  int    `json:"audio_duration,omitempty"`
	InferenceSteps *int   `json:"inference_steps,omitempty"`
	Seed           *int64 `json:"seed,omitempty"`
}

// queryItem is one task's poll state. `result` arrives as a JSON-encoded
// STRING (the ACE-Step API double-encodes it).
type queryItem struct {
	TaskID string `json:"task_id"`
	Status int    `json:"status"` // 0 queued/running, 1 succeeded, 2 failed
	Result string `json:"result"`
}

// 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.
func (m *musicModel) Generate(ctx context.Context, req musicgen.Request, opts ...musicgen.Option) (*musicgen.Result, error) {
	req = req.Apply(opts...)
	if strings.TrimSpace(req.Prompt) == "" {
		return nil, fmt.Errorf("%w: music generation requires a prompt", llm.ErrUnsupported)
	}
	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 {
		return nil, err
	}
	item, err := m.pollResult(ctx, taskID)
	if err != nil {
		return nil, err
	}
	return m.fetchResult(ctx, req.Format, item)
}

// releaseTask submits the job and returns its task id.
func (m *musicModel) releaseTask(ctx context.Context, req musicgen.Request) (string, error) {
	path, err := upstreamPath(m.id, "/release_task")
	if err != nil {
		return "", err
	}
	wire := releaseTaskRequest{
		Prompt:         req.Prompt,
		Lyrics:         req.Lyrics,
		AudioFormat:    req.Format,
		TaskType:       "text2music",
		AudioDuration:  req.DurationSeconds,
		InferenceSteps: req.Steps,
		Seed:           req.Seed,
	}
	// Tolerant envelope: {"data": {"task_id": ...}} per the docs, with a
	// top-level fallback in case the wrapper changes.
	var resp struct {
		Data struct {
			TaskID string `json:"task_id"`
		} `json:"data"`
		TaskID string `json:"task_id"`
	}
	if err := m.p.doJSON(ctx, http.MethodPost, path, m.id, &wire, &resp); err != nil {
		return "", err
	}
	taskID := resp.Data.TaskID
	if taskID == "" {
		taskID = resp.TaskID
	}
	if taskID == "" {
		return "", &llm.APIError{Provider: m.p.name, Model: m.id, Message: "release_task returned no task_id"}
	}
	return taskID, nil
}

// pollResult polls query_result until the task succeeds, fails, or ctx
// 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 {
		return nil, err
	}
	body := map[string]any{"task_id_list": []string{taskID}}
	ticker := time.NewTicker(musicPollInterval)
	defer ticker.Stop()
	failures := 0
	var lastErr error
	for {
		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 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():
			return nil, fmt.Errorf("llama-swap: music generation: %w", ctx.Err())
		case <-ticker.C:
		}
	}
}

// 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) {
	var env struct {
		Data json.RawMessage `json:"data"`
	}
	candidates := raw
	if json.Unmarshal(raw, &env) == nil && len(env.Data) > 0 {
		candidates = env.Data
	}
	var items []queryItem
	if err := json.Unmarshal(candidates, &items); err != nil {
		var one queryItem
		if err := json.Unmarshal(candidates, &one); err != nil {
			return nil, fmt.Errorf("unrecognized query_result payload shape")
		}
		items = []queryItem{one}
	}
	for i := range items {
		// 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
		}
	}
	return nil, fmt.Errorf("query_result did not include task %s", taskID)
}

// fetchResult downloads the finished clip named by the job's result blob.
func (m *musicModel) fetchResult(ctx context.Context, format string, item *queryItem) (*musicgen.Result, error) {
	var result struct {
		File string `json:"file"`
	}
	if err := json.Unmarshal([]byte(item.Result), &result); err != nil || result.File == "" {
		return nil, &llm.APIError{Provider: m.p.name, Model: m.id,
			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. 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
	}
	raw, contentType, err := m.p.doRaw(ctx, http.MethodGet, path, m.id, "", nil, maxResponseBytes)
	if err != nil {
		return nil, err
	}
	if len(raw) == 0 {
		return nil, &llm.APIError{Provider: m.p.name, Model: m.id, Message: "music response contained no audio"}
	}
	mimeType := musicFormatMIME(contentType, format)
	return &musicgen.Result{
		Audio: musicgen.Audio{Data: raw, MIME: mimeType},
		Raw:   json.RawMessage(item.Result),
	}, nil
}