decode.go

/*
Copyright 2026 The llm-d Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package steps

import (
	"context"
	"errors"
	"fmt"
	"net/http"

	"sigs.k8s.io/controller-runtime/pkg/log"

	logutil "github.com/llm-d/llm-d-router/pkg/common/observability/logging"
	reqcommon "github.com/llm-d/llm-d-router/pkg/common/request"

	"github.com/llm-d/llm-d-router/pkg/coordinator/connectors/kv"
	"github.com/llm-d/llm-d-router/pkg/coordinator/gateway"
	coordmetrics "github.com/llm-d/llm-d-router/pkg/coordinator/metrics"
	"github.com/llm-d/llm-d-router/pkg/coordinator/pipeline"
)

const DecodeStepName = "decode"

func init() {
	pipeline.Register(DecodeStepName, NewDecodeStep)
}

type DecodeStep struct {
	useOpenAIFormat bool
	gwClient        *gateway.Client
	kv              kv.Connector
}

func NewDecodeStep(gwClient *gateway.Client, params map[string]any) (pipeline.Step, error) {
	if gwClient == nil {
		return nil, errors.New("decode: gateway client is required")
	}
	useOpenAI, err := parseUseOpenAIFormat(params)
	if err != nil {
		return nil, fmt.Errorf("decode: %w", err)
	}
	kvName, err := paramString(params, ParamKVConnector)
	if err != nil {
		return nil, fmt.Errorf("decode: %w", err)
	}
	kvConn, err := kv.Build(kvName)
	if err != nil {
		return nil, fmt.Errorf("decode: %w", err)
	}
	return &DecodeStep{useOpenAIFormat: useOpenAI, gwClient: gwClient, kv: kvConn}, nil
}

func (s *DecodeStep) Name() string { return DecodeStepName }

func (s *DecodeStep) Execute(ctx context.Context, reqCtx *pipeline.RequestContext) error {
	logger := log.FromContext(ctx).WithName(DecodeStepName)

	s.prepareDecodeBody(ctx, reqCtx)

	logger.V(logutil.DEFAULT).Info("sending request", "path", reqCtx.OriginalPath, "stream", reqCtx.Stream)

	proxyReq, err := newDecodeProxyRequest(ctx, logger, DecodeStepName, reqCtx, s.gwClient, reqCtx.Body, nil)
	if err != nil {
		return err
	}

	transport := instrumentedTransport(s.gwClient.Transport(), coordmetrics.UpstreamDecode)
	proxy, out := newDecodeProxy(logger, transport, nil)
	proxy.ServeHTTP(reqCtx.ResponseWriter, proxyReq)
	if out.TransportErr != nil {
		return &pipeline.UpstreamStreamedError{Step: DecodeStepName, Cause: out.TransportErr}
	}
	if out.Status >= http.StatusBadRequest {
		return &pipeline.UpstreamStreamedError{Step: DecodeStepName, StatusCode: out.Status}
	}
	return nil
}

// prepareDecodeBody mutates reqCtx.Body in place rather than on a clone (unlike
// prefill and conditional-decode). decode is the terminal pipeline step: its body
// is streamed straight to the client and no later step reads reqCtx.Body. A clone
// would also be insufficient, since injectUUIDs mutates nested values that a shallow
// maps.Clone would still share. This is sound only while the pipeline runs steps
// sequentially; if it ever goes concurrent, decode must copy like the others.
func (s *DecodeStep) prepareDecodeBody(ctx context.Context, reqCtx *pipeline.RequestContext) {
	kvParams := s.kv.PrepareDecodeKVParams(ctx, reqCtx)
	s.injectUUIDs(reqCtx)

	format := resolveFormat(s.useOpenAIFormat, reqCtx.OriginalPath)
	switch format {
	case gateway.FormatChatCompletions:
		reqCtx.Body[reqcommon.FieldKVTransferParams] = kvParams
		s.injectTokensField(reqCtx)
	case gateway.FormatCompletions:
		reqCtx.Body[reqcommon.FieldKVTransferParams] = kvParams
		if len(reqCtx.TokenIDs) > 0 {
			reqCtx.Body["prompt"] = reqCtx.TokenIDs
		}
	case gateway.FormatGenerate:
		// The /inference/v1/generate engine reads transfer params only from
		// sampling_params.extra_args; a top-level kv_transfer_params is ignored,
		// so the decode worker never pulls the prefill KV over NIXL. Merge into
		// the client's sampling_params to preserve max_tokens and other fields.
		sampling, ok := reqCtx.Body[reqcommon.FieldSamplingParams].(map[string]any)
		if !ok {
			sampling = map[string]any{}
			reqCtx.Body[reqcommon.FieldSamplingParams] = sampling
		}
		setGenerateTransferParams(sampling, kvParams, nil)
	}
}

func (s *DecodeStep) injectTokensField(reqCtx *pipeline.RequestContext) {
	tokens := map[string]any{
		"token_ids": reqCtx.TokenIDs,
	}
	if features := buildMMFeatures(reqCtx.MultimodalEntries, false); features != nil {
		tokens["features"] = features
	}
	reqCtx.Body["tokens"] = tokens
}

// injectUUIDs tags each media content part with the uuid the decode
// backend uses for prefix-cache keying. Each part is paired with the
// MultimodalEntry that shares its modality and local index — the same
// invariant the encode fanout uses (see localIndexOf in utils.go).
// Non-media parts (text, tool_use, unknown types) are skipped.
func (s *DecodeStep) injectUUIDs(reqCtx *pipeline.RequestContext) {
	messages, ok := reqCtx.Body["messages"].([]any)
	if !ok {
		return
	}

	// Build a (modality, localIndex) → Hash lookup once so per-part
	// tagging stays O(1) even on requests with many media items.
	hashByModLocalIdx := make(map[string]map[int]string)
	for i, entry := range reqCtx.MultimodalEntries {
		mod := entryModality(entry)
		localIdx := localIndexOf(reqCtx.MultimodalEntries, i)
		if _, ok := hashByModLocalIdx[mod]; !ok {
			hashByModLocalIdx[mod] = make(map[int]string)
		}
		hashByModLocalIdx[mod][localIdx] = entry.Hash
	}

	// Walk parts, incrementing a per-modality counter so each media part
	// gets the hash of the entry at the matching (modality, localIndex).
	modCounter := make(map[string]int)
	for _, msg := range messages {
		msgMap, ok := msg.(map[string]any)
		if !ok {
			continue
		}
		content, ok := msgMap["content"].([]any)
		if !ok {
			continue
		}
		for _, part := range content {
			partMap, ok := part.(map[string]any)
			if !ok {
				continue
			}
			partType, _ := partMap["type"].(string)
			modality, isMedia := partTypeModality[partType]
			if !isMedia {
				continue
			}
			localIdx := modCounter[modality]
			modCounter[modality]++
			if h, ok := hashByModLocalIdx[modality][localIdx]; ok {
				partMap["uuid"] = h
			}
		}
	}
}