DeepSeek-Reasonix/internal/runtime/agent/sampling_request.go

package agent

import (
	"context"
	"encoding/json"
	"fmt"
	"time"

	"reasonix/internal/contract/event"
	"reasonix/internal/contract/provider"
)

// samplingRequest is a once-prepared, frozen provider request for one model
// round. All stream retries replay this exact payload — no synthetic recovery
// messages, no schema reorder, no previous_response_id drift from failed attempts.
type samplingRequest struct {
	req provider.Request
	// overflowFolded records that this round already answered a context-overflow
	// rejection by folding. The recovery is one-shot: a second rejection means
	// the fold did not reach far enough, not that another one will.
	overflowFolded bool
	// overflowRecovery is why that fold did not produce a smaller request, kept
	// so the failure the user sees names the recovery that was tried.
	overflowRecovery error
	// sent is what the visible view held when this request was built, the basis
	// the provider's reported size is later anchored to.
	sent usageAnchor
}

func (a *Agent) streamProviderRequest(ctx context.Context, req provider.Request) (<-chan provider.Chunk, error) {
	ch, err := a.svc.prov.Stream(ctx, req)
	if err == nil {
		a.window().rememberSummaryPrefix(req)
	}
	return ch, err
}

func (a *Agent) handleSamplingError(
	ctx context.Context,
	attemptID string,
	attempt int,
	frozen *samplingRequest,
	result, last streamedTurn,
	billable *provider.Usage,
) (retry bool, terminal streamedTurn) {
	if provider.IsStreamInterrupted(result.err) && attempt < maxSamplingAttempts {
		reason := provider.StreamInterruptReason(result.err)
		a.emitStreamAttempt(attemptID, event.StreamAttemptDiscard, attempt, reason, result.err, 0)
		delay := streamRetryDelay(attempt)
		a.svc.sink.Emit(event.Event{
			Kind: event.Retrying, RetryAttempt: attempt, RetryMax: maxStreamRecoveries,
			RetryScope: event.RetryScopeStream, RetryCause: provider.RetryCauseOfStreamInterrupt(reason),
			RetryDelayMs: delay.Milliseconds(),
		})
		if !streamRetrySleep(ctx, delay) {
			return false, streamedTurn{usage: finalizeSamplingUsage(billable, result.usage), attemptID: attemptID, interrupted: true, err: ctx.Err()}
		}
		return true, streamedTurn{}
	}
	if attempt < maxSamplingAttempts && a.window().recoverContextOverflow(ctx, frozen, result.err) {
		a.emitStreamAttempt(attemptID, event.StreamAttemptDiscard, attempt, "context overflow", result.err, 0)
		return true, streamedTurn{}
	}
	// Exhausted retries or non-retryable error: leave the last speculative UI
	// visible (no discard) so LocalOnly can mirror it.
	if frozen != nil && frozen.overflowRecovery != nil && last.err != nil {
		last.err = fmt.Errorf("%w (context recovery did not help: %w)", last.err, frozen.overflowRecovery)
	}
	last.usage = finalizeSamplingUsage(billable, result.usage)
	return false, last
}

// prepareSamplingRequest freezes one model-round request (preflight + interceptors).
// Output budgets are resolved only here and never change the compact_ratio
// trigger. Physical overflow may attempt at most one recovery summary.
func (a *Agent) prepareSamplingRequest(ctx context.Context) (samplingRequest, error) {
	frozen, err := a.buildSamplingRequest(ctx, CompactionTriggerPressure)
	if err != nil {
		return samplingRequest{}, err
	}
	if budget, clipped, budgetErr := a.window().effectiveOutputBudget(frozen.req); budgetErr != nil {
		// One-shot physical overflow recovery. Do not loop.
		if _, perr := a.window().contextManager().Prepare(ctx, ContextPreparePolicy{
			Trigger: CompactionTriggerOverflow,
		}); perr != nil {
			return samplingRequest{}, budgetErr
		}
		rebuilt, rerr := a.buildSamplingRequest(ctx, CompactionTriggerPressure)
		if rerr != nil {
			return samplingRequest{}, rerr
		}
		if _, _, budgetErr2 := a.window().effectiveOutputBudget(rebuilt.req); budgetErr2 != nil {
			return samplingRequest{}, budgetErr2
		}
		// Re-apply clipping on the recovered view.
		if budget2, clipped2, err2 := a.window().effectiveOutputBudget(rebuilt.req); err2 == nil && clipped2 {
			rebuilt.req.MaxTokens = budget2
		}
		shape := a.window().requestCalibrationShape(rebuilt.req)
		a.sess.output.activeReqShape.Store(&shape)
		a.sess.output.usageAnchor.Store(&rebuilt.sent)
		return samplingRequest{req: freezeProviderRequest(rebuilt.req), sent: rebuilt.sent}, nil
	} else if clipped {
		frozen.req.MaxTokens = budget
	}
	shape := a.window().requestCalibrationShape(frozen.req)
	a.sess.output.activeReqShape.Store(&shape)
	a.sess.output.usageAnchor.Store(&frozen.sent)
	return samplingRequest{req: freezeProviderRequest(frozen.req), sent: frozen.sent}, nil
}

func (a *Agent) buildSamplingRequest(ctx context.Context, trigger string) (samplingRequest, error) {
	// CreatedAt is durable UI metadata, not model input. Strip it from the
	// transport copy so wall-clock differences never invalidate the provider's
	// prompt-cache prefix (and custom providers cannot accidentally send it).
	prepared, err := a.window().contextManager().Prepare(ctx, ContextPreparePolicy{Trigger: trigger})
	if err != nil {
		return samplingRequest{}, err
	}
	// The host tails are already here: Prepare answers from modelVisibleMessages,
	// and the host context tail appends whatever it is given again.
	requestMessages := append([]provider.Message(nil), provider.ModelMessages(prepared.Messages)...)
	requestMessages = a.providerProjectionMessages(requestMessages)
	for i := range requestMessages {
		requestMessages[i].CreatedAt = 0
	}
	// context.prepare: extensions may rewrite the message copy feeding THIS
	// request. The session log is never touched — the replacement is
	// ephemeral, so the next request starts from the unmodified history.
	requestMessages, err = a.interceptContextPrepare(ctx, requestMessages)
	if err != nil {
		return samplingRequest{}, err
	}
	req := provider.Request{
		Messages:       requestMessages,
		Tools:          a.svc.tools.ProviderSchemas(ctx),
		MaxTokens:      a.maxOutputTokens,
		Temperature:    provider.OptionalTemperature(a.temperature),
		ResponseFormat: responseFormatFromRequest(ctx),
	}
	// provider.request: the fully assembled request gets one last ruling
	// (revalidated by the payload registry) before it goes on the wire.
	req, err = a.interceptProviderRequest(ctx, req)
	if err != nil {
		return samplingRequest{}, err
	}
	// Host-owned and set after the extension ruling: the payload it rewrites
	// has no mode, and a replaced request must not drop the session's choice.
	req.Mode = a.sess.mode.get()
	// The same rule for the session's switched depth; an unset one leaves any override an extension set.
	if level := a.sess.effort.get(); level != "" {
		req.EffortOverride = level
	}
	sent := usageAnchor{
		visible:           a.window().visibleRequestShape(prepared.Messages),
		projectionVersion: prepared.ProjectionVersion,
	}
	return samplingRequest{req: req, sent: sent}, nil
}

// providerProjectionMessages applies provider-specific role compatibility to a
// request copy. Projection sidecars retain logical user-turn boundaries so
// explicit range compression can continue to resolve anchors across calls.
func (a *Agent) providerProjectionMessages(msgs []provider.Message) []provider.Message {
	if a != nil && a.strictAlternatingRoles {
		return coalesceProjectionUserRuns(msgs)
	}
	return msgs
}

// freezeProviderRequest deep-copies the provider-visible request surface so
// retries share identical messages, tools order, temperature, and format.
func freezeProviderRequest(req provider.Request) provider.Request {
	out := req
	if len(req.Messages) > 0 {
		out.Messages = append([]provider.Message(nil), req.Messages...)
		for i := range out.Messages {
			if len(out.Messages[i].ToolCalls) > 0 {
				out.Messages[i].ToolCalls = append([]provider.ToolCall(nil), out.Messages[i].ToolCalls...)
			}
			if len(out.Messages[i].Images) > 0 {
				out.Messages[i].Images = append([]string(nil), out.Messages[i].Images...)
			}
			if len(out.Messages[i].ResponsesItems) > 0 {
				items := make([]json.RawMessage, len(out.Messages[i].ResponsesItems))
				for j, item := range out.Messages[i].ResponsesItems {
					items[j] = append(json.RawMessage(nil), item...)
				}
				out.Messages[i].ResponsesItems = items
			}
		}
	}
	if len(req.Tools) > 0 {
		out.Tools = make([]provider.ToolSchema, len(req.Tools))
		for i, schema := range req.Tools {
			out.Tools[i] = schema
			if len(schema.Parameters) > 0 {
				out.Tools[i].Parameters = append(json.RawMessage(nil), schema.Parameters...)
			}
		}
	}
	if req.Temperature != nil {
		t := *req.Temperature
		out.Temperature = &t
	}
	if req.ResponseFormat != nil {
		rf := *req.ResponseFormat
		out.ResponseFormat = &rf
	}
	return out
}

// withHeaderRetries makes a request report its connection retries to the sink
// and, when the user set one, apply their retry limit.
func withHeaderRetries(ctx context.Context, sink event.Sink, limit int) context.Context {
	ctx = provider.WithRetryNotify(ctx, headerRetryNotice(sink))
	if limit != 0 {
		ctx = provider.WithRetryLimit(ctx, limit)
	}
	return ctx
}

// headerRetryNotice turns a connection+header backoff into the Retrying event a
// frontend words, carrying what the provider knew about why it failed.
func headerRetryNotice(sink event.Sink) provider.RetryNotify {
	return func(info provider.RetryInfo) {
		sink.Emit(event.Event{
			Kind: event.Retrying, RetryAttempt: info.Attempt, RetryMax: info.Max, RetryScope: event.RetryScopeHeaders,
			RetryCause: info.Cause, RetryStatus: info.Status,
			RetryDelayMs: info.Delay.Milliseconds(), RetryTimeoutSecs: int(info.Timeout / time.Second),
		})
	}
}