iop/apps/edge/internal/openai/provider_model_rewrite.go

package openai

import (
	"bytes"
	"encoding/json"
	"fmt"
	"strings"
	"sync"
)

type openAIProviderBodyBuilder struct {
	mu     sync.Mutex
	build  func(string) (*openAIRebuiltLease, error)
	lease  *openAIRebuiltLease
	built  bool
	closed bool
}

func newOpenAIProviderBodyBuilder(build func(string) (*openAIRebuiltLease, error)) *openAIProviderBodyBuilder {
	return &openAIProviderBodyBuilder{build: build}
}

func (b *openAIProviderBodyBuilder) BuildBody(target string) ([]byte, error) {
	if b == nil {
		return nil, fmt.Errorf("OpenAI provider body builder is unavailable")
	}
	b.mu.Lock()
	if b.closed || b.built || b.build == nil {
		b.mu.Unlock()
		return nil, fmt.Errorf("OpenAI provider body builder is unavailable")
	}
	b.built = true
	build := b.build
	b.mu.Unlock()

	lease, err := build(target)
	if err != nil {
		return nil, err
	}
	b.mu.Lock()
	if b.closed {
		b.mu.Unlock()
		lease.release()
		return nil, fmt.Errorf("OpenAI provider body builder is unavailable")
	}
	b.lease = lease
	b.mu.Unlock()
	body, err := lease.body()
	if err != nil {
		b.Close()
		return nil, err
	}
	return body, nil
}

func (b *openAIProviderBodyBuilder) Close() {
	if b == nil {
		return
	}
	b.mu.Lock()
	b.closed = true
	b.build = nil
	lease := b.lease
	b.lease = nil
	b.mu.Unlock()
	lease.release()
}

type providerModelRewriter struct {
	streaming bool
	model     string
	pending   []byte
}

func newProviderModelRewriter(streaming bool, model string) *providerModelRewriter {
	model = strings.TrimSpace(model)
	if model == "" {
		return nil
	}
	return &providerModelRewriter{streaming: streaming, model: model}
}

func (r *providerModelRewriter) AppendStream(chunk []byte) []byte {
	if r == nil || !r.streaming || len(chunk) == 0 {
		return chunk
	}
	r.pending = append(r.pending, chunk...)
	var out bytes.Buffer
	for {
		idx := bytes.IndexByte(r.pending, '\n')
		if idx < 0 {
			break
		}
		line := r.pending[:idx+1]
		out.Write(rewriteProviderSSEModelLine(line, r.model))
		r.pending = r.pending[idx+1:]
	}
	return out.Bytes()
}

func (r *providerModelRewriter) FlushStream() []byte {
	if r == nil || len(r.pending) == 0 {
		return nil
	}
	pending := r.pending
	r.pending = nil
	return rewriteProviderSSEModelLine(pending, r.model)
}

func (r *providerModelRewriter) RewriteComplete(body []byte) []byte {
	if r == nil || len(body) == 0 {
		return body
	}
	return rewriteProviderJSONModel(body, r.model)
}

func rewriteProviderSSEModelLine(line []byte, model string) []byte {
	body, ending := splitLineEnding(line)
	prefix, payload, ok := bytes.Cut(body, []byte(":"))
	if !ok || strings.TrimSpace(string(prefix)) != "data" {
		return line
	}
	payload = bytes.TrimSpace(payload)
	if len(payload) == 0 || bytes.Equal(payload, []byte("[DONE]")) {
		return line
	}
	rewritten := rewriteProviderJSONModel(payload, model)
	if bytes.Equal(rewritten, payload) {
		return line
	}
	out := make([]byte, 0, len("data: ")+len(rewritten)+len(ending))
	out = append(out, "data: "...)
	out = append(out, rewritten...)
	out = append(out, ending...)
	return out
}

func splitLineEnding(line []byte) ([]byte, []byte) {
	if len(line) == 0 || line[len(line)-1] != '\n' {
		return line, nil
	}
	if len(line) >= 2 && line[len(line)-2] == '\r' {
		return line[:len(line)-2], []byte("\r\n")
	}
	return line[:len(line)-1], []byte("\n")
}

func rewriteProviderJSONModel(body []byte, model string) []byte {
	modelJSON, err := json.Marshal(model)
	if err != nil {
		return body
	}
	fields, _, err := scanTopLevelJSONObject(body)
	if err != nil {
		return body
	}
	for _, field := range fields {
		if field.name == "model" {
			plan, err := planTopLevelJSONPatches(body, []topLevelJSONPatch{{name: "model", value: modelJSON}})
			if err == nil {
				return plan.apply()
			}
			break
		}
	}
	return body
}

// rewriteChatCompletionModel replaces only the model field of the caller's
// original request JSON so the provider receives its served model name. The
// rest of the caller payload is forwarded without IOP rewriting.
func rewriteChatCompletionModel(rawBody []byte, target string, req chatCompletionRequest) ([]byte, error) {
	patches, err := chatCompletionProviderPatches(target, req)
	if err != nil {
		return nil, err
	}
	if len(patches) == 0 {
		return rawBody, nil
	}
	plan, err := planTopLevelJSONPatches(rawBody, patches)
	if err != nil {
		return nil, err
	}
	return plan.apply(), nil
}

func rewriteChatCompletionModelFromIngress(ingress *openAIIngressSnapshot, target string, req chatCompletionRequest) (*openAIRebuiltLease, error) {
	patches, err := chatCompletionProviderPatches(target, req)
	if err != nil {
		return nil, err
	}
	return rewriteIngressTopLevelJSON(ingress, patches)
}

func chatCompletionProviderPatches(target string, req chatCompletionRequest) ([]topLevelJSONPatch, error) {
	patches := make([]topLevelJSONPatch, 0, 5)
	if strings.TrimSpace(target) != "" {
		modelJSON, err := json.Marshal(target)
		if err != nil {
			return nil, err
		}
		patches = append(patches, topLevelJSONPatch{name: "model", value: modelJSON})
	}
	if req.MaxTokens != nil {
		maxTokensJSON, err := json.Marshal(*req.MaxTokens)
		if err != nil {
			return nil, err
		}
		patches = append(patches,
			topLevelJSONPatch{name: "max_tokens", value: maxTokensJSON},
			topLevelJSONPatch{name: "max_completion_tokens"},
		)
	}
	if req.ThinkingTokenBudget != nil {
		budgetJSON, err := json.Marshal(*req.ThinkingTokenBudget)
		if err != nil {
			return nil, err
		}
		patches = append(patches, topLevelJSONPatch{name: "thinking_token_budget", value: budgetJSON})
	}
	if req.Think != nil {
		thinkJSON, err := json.Marshal(*req.Think)
		if err != nil {
			return nil, err
		}
		patches = append(patches, topLevelJSONPatch{name: "think", value: thinkJSON})
	}
	return patches, nil
}

// rewriteResponsesModel replaces only the model field of the caller's original
// /v1/responses request JSON so the provider receives its served model name.
// Every other field (input, instructions, tools, max_output_tokens, and any
// Codex/Responses-specific field) is forwarded without IOP rewriting. An empty
// target leaves the body untouched; invalid JSON is rejected.
func rewriteResponsesModel(rawBody []byte, target string) ([]byte, error) {
	if strings.TrimSpace(target) == "" {
		return rawBody, nil
	}
	modelJSON, err := json.Marshal(target)
	if err != nil {
		return nil, err
	}
	plan, err := planTopLevelJSONPatches(rawBody, []topLevelJSONPatch{{name: "model", value: modelJSON}})
	if err != nil {
		return nil, err
	}
	return plan.apply(), nil
}

func rewriteResponsesModelFromIngress(ingress *openAIIngressSnapshot, target string) (*openAIRebuiltLease, error) {
	if strings.TrimSpace(target) == "" {
		return rewriteIngressTopLevelJSON(ingress, nil)
	}
	modelJSON, err := json.Marshal(target)
	if err != nil {
		return nil, err
	}
	return rewriteIngressTopLevelJSON(ingress, []topLevelJSONPatch{{name: "model", value: modelJSON}})
}

// rewriteIngressTopLevelJSON places both allocation checks around the single
// patched output and returns a lease that stays live through synchronous
// provider submission. The canonical ingress owner remains unchanged.
func rewriteIngressTopLevelJSON(ingress *openAIIngressSnapshot, patches []topLevelJSONPatch) (*openAIRebuiltLease, error) {
	body, err := ingress.canonicalBody()
	if err != nil {
		return nil, err
	}
	plan, err := planTopLevelJSONPatches(body, patches)
	if err != nil {
		return nil, err
	}
	if len(plan.edits) == 0 {
		return &openAIRebuiltLease{ingress: ingress, bodyAlias: body}, nil
	}
	guard, err := ingress.reserveRebuild(int64(plan.outputSize))
	if err != nil {
		return nil, err
	}
	output := plan.apply()
	rebuilt, err := guard.CommitOwnedTyped(openAIRebuiltBodyViewName, output)
	if err != nil {
		guard.Close()
		return nil, err
	}
	return &openAIRebuiltLease{
		ingress: ingress, guard: guard, rebuilt: rebuilt, bodyAlias: output,
	}, nil
}