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

package openai

import (
	"context"
	"fmt"
	"net/http"
	"strings"
)

type hotPathEndpointError struct {
	Status  int
	Type    string
	Message string
}

type hotPathTerminalIntent struct {
	Output normalizedStageOutput
	Error  *hotPathEndpointError
}

type hotPathCleanupTurn struct {
	RequestID string
	Output    normalizedStageOutput
}

func (i hotPathTerminalIntent) clone() hotPathTerminalIntent {
	out := hotPathTerminalIntent{Output: cloneNormalizedStageOutput(i.Output)}
	if i.Error != nil {
		endpointErr := *i.Error
		out.Error = &endpointErr
	}
	return out
}

func (i hotPathTerminalIntent) terminalClass() string {
	if i.Error != nil {
		return "primary_error"
	}
	return "success"
}

func (s *hotPathLightStore) beginCleanup(
	ctx context.Context,
	requestID, ownerEdgeID string,
	intent hotPathTerminalIntent,
	coordinator *logicalRequestCoordinator,
) (normalizedStageOutput, error) {
	if s == nil || coordinator == nil {
		return normalizedStageOutput{}, fmt.Errorf("light cleanup is unavailable")
	}
	s.mu.Lock()
	defer s.mu.Unlock()
	record := s.records[requestID]
	if record == nil || record.ownerEdgeID != ownerEdgeID || !record.running || record.pending != nil {
		return normalizedStageOutput{}, fmt.Errorf("review completion cannot enter cleanup")
	}
	if record.phase != hotPathPhaseReviewResolution && record.phase != hotPathPhaseReviewRepair {
		return normalizedStageOutput{}, fmt.Errorf("review completion is not resolution or repair")
	}
	return s.beginCleanupLocked(ctx, record, record.reviewStageID, intent, coordinator)
}

func (s *hotPathLightStore) beginPrimaryErrorCleanup(
	ctx context.Context,
	requestID, ownerEdgeID string,
	primary hotPathEndpointError,
	coordinator *logicalRequestCoordinator,
) (normalizedStageOutput, error) {
	if s == nil || coordinator == nil {
		return normalizedStageOutput{}, fmt.Errorf("light cleanup is unavailable")
	}
	s.mu.Lock()
	defer s.mu.Unlock()
	record := s.records[requestID]
	if record == nil || record.ownerEdgeID != ownerEdgeID || record.cleanupTransitions != 0 || record.terminalIntent != nil {
		return normalizedStageOutput{}, fmt.Errorf("primary-error cleanup is unavailable")
	}
	fromStageID, err := record.primaryErrorCleanupSource()
	if err != nil {
		return normalizedStageOutput{}, err
	}
	intent := hotPathTerminalIntent{Error: &primary}
	return s.beginCleanupLocked(ctx, record, fromStageID, intent, coordinator)
}

func (r *hotPathLightRecord) primaryErrorCleanupSource() (string, error) {
	if r == nil || r.running || r.pending != nil || r.cleanupTransitions != 0 || r.terminalIntent != nil {
		return "", fmt.Errorf("primary-error cleanup source is unavailable")
	}
	if r.selectorCommit.StageID != r.selectorStageID || strings.TrimSpace(r.selectorCommit.ResponseID) == "" {
		return "", fmt.Errorf("primary-error cleanup selector correlation is unavailable")
	}

	switch r.phase {
	case hotPathPhaseAwaitArtifacts:
		if r.localStageID != "" || r.reviewStageID != "" {
			return "", fmt.Errorf("primary-error cleanup artifact source is mismatched")
		}
		return "", nil
	case hotPathPhaseLocalActive:
		if !r.artifactReady || !validLogicalRequestID(r.localStageID) || r.reviewStageID != "" {
			return "", fmt.Errorf("primary-error cleanup local source is mismatched")
		}
		return r.localStageID, nil
	case hotPathPhaseReviewActive, hotPathPhaseReviewAwaitRead, hotPathPhaseReviewResolution, hotPathPhaseReviewRepair:
		if !r.artifactReady || !validLogicalRequestID(r.localStageID) || !validLogicalRequestID(r.reviewStageID) ||
			r.localCommit.StageID != r.localStageID || strings.TrimSpace(r.localCommit.ResponseID) == "" {
			return "", fmt.Errorf("primary-error cleanup review source is mismatched")
		}
		return r.reviewStageID, nil
	default:
		return "", fmt.Errorf("phase %q cannot enter primary-error cleanup", r.phase)
	}
}

func (s *hotPathLightStore) beginCleanupLocked(
	ctx context.Context,
	record *hotPathLightRecord,
	fromStageID string,
	intent hotPathTerminalIntent,
	coordinator *logicalRequestCoordinator,
) (normalizedStageOutput, error) {
	if err := ctx.Err(); err != nil {
		record.running = false
		_ = coordinator.disconnect(record.requestID, record.ownerEdgeID, "cancelled")
		return normalizedStageOutput{}, err
	}
	if record.cleanupTransitions != 0 || record.terminalIntent != nil {
		return normalizedStageOutput{}, fmt.Errorf("cleanup pending was already committed")
	}

	cleanupStageID, err := coordinator.newStageID()
	if err != nil {
		return normalizedStageOutput{}, err
	}
	providerCallID, err := coordinator.newCallID()
	if err != nil {
		return normalizedStageOutput{}, err
	}
	paths := newReservedPaths(record.requestID)
	deleteBinding := record.binding.operation(opKindDelete)
	if deleteBinding == nil {
		return normalizedStageOutput{}, fmt.Errorf("cleanup delete binding is unavailable")
	}
	deleteArgs := make(map[string]any)
	setMappedArgument(deleteArgs, deleteBinding.pathField, paths.JobDir)
	providerCall := normalizedToolCall{
		ID: providerCallID, ProviderCallID: providerCallID, Name: deleteBinding.toolName,
		Arguments: deleteArgs, Path: paths.JobDir,
	}
	mapped, payload, err := mapArtifactCall(record.binding, providerCall, opKindDelete, paths.JobDir, coordinator)
	if err != nil {
		return normalizedStageOutput{}, fmt.Errorf("map cleanup delete: %w", err)
	}
	responseID := strings.TrimSpace(intent.Output.ResponseID)
	if responseID == "" {
		responseID = strings.TrimSpace(record.selectorCommit.ResponseID)
	}
	if responseID == "" {
		return normalizedStageOutput{}, fmt.Errorf("cleanup response identity is unavailable")
	}
	cleanupOutput := normalizedStageOutput{
		ResponseID: responseID, Created: intent.Output.Created,
		ToolCalls: []normalizedToolCall{mapped}, TerminalReason: "tool_calls",
	}
	if record.protocol == "anthropic" {
		cleanupOutput.TerminalReason = "tool_use"
	}
	issuedHash, err := directIssuedCallHash(record.protocol, cleanupOutput)
	if err != nil {
		return normalizedStageOutput{}, fmt.Errorf("fingerprint cleanup call: %w", err)
	}
	if _, err := coordinator.startCleanup(record.requestID, record.ownerEdgeID, fromStageID, cleanupStageID, intent.terminalClass()); err != nil {
		return normalizedStageOutput{}, err
	}
	if _, err := coordinator.awaitToolResults(record.requestID, record.ownerEdgeID, cleanupStageID, []logicalRequestExpectedTool{{
		PublicCallID: mapped.ID, ProviderCallID: mapped.ProviderCallID,
	}}, issuedHash); err != nil {
		return normalizedStageOutput{}, err
	}

	stored := intent.clone()
	record.terminalIntent = &stored
	record.pendingKind = hotPathPendingCleanup
	record.pending = map[string]hotPathPendingCall{
		mapped.ID: {publicCallID: mapped.ID, providerCallID: mapped.ProviderCallID, payload: payload},
	}
	record.pendingHash = issuedHash
	record.pendingOutput = cloneNormalizedStageOutput(cleanupOutput)
	record.phase = hotPathPhaseCleanupPending
	record.cleanupTransitions++
	record.running = false
	return cleanupOutput, nil
}

func (s *hotPathLightStore) consumeCleanupLocked(
	record *hotPathLightRecord,
	lineage logicalRequestContinuationLineage,
	results []workspaceResult,
	coordinator *logicalRequestCoordinator,
) (logicalRequestSnapshot, hotPathLightDisposition, bool, error) {
	if record.terminalIntent == nil || len(record.pending) != 1 || len(results) != 1 {
		return logicalRequestSnapshot{}, hotPathLightDisposition{}, true, fmt.Errorf("cleanup result set mismatch")
	}
	result := results[0]
	pending, ok := record.pending[result.callID]
	if !ok || pending.payload == nil {
		return logicalRequestSnapshot{}, hotPathLightDisposition{}, true, fmt.Errorf("cleanup result id is not pending")
	}
	if reason := matchResultCorrelation(record.binding, pending.payload, result); reason != "" {
		return logicalRequestSnapshot{}, hotPathLightDisposition{}, true, fmt.Errorf("cleanup receipt rejected: %s", reason)
	}

	intent := record.terminalIntent.clone()
	receipt := matchResultReceipt(record.binding, pending.payload, result)
	if !receipt.matched && intent.Error == nil {
		intent.Error = standardCleanupEndpointError(record.protocol)
		intent.Output = normalizedStageOutput{}
	}
	snap, err := coordinator.commitCleanupByLineage(record.ownerEdgeID, record.principalRef, lineage)
	if err != nil {
		return logicalRequestSnapshot{}, hotPathLightDisposition{}, true, err
	}
	requestID := record.requestID
	stageID := snap.ActiveStageID
	delete(s.records, requestID)
	return snap, hotPathLightDisposition{
		RequestID: requestID, StageID: stageID, Phase: hotPathPhaseCleanupPending, Terminal: &intent,
	}, true, nil
}

func standardCleanupEndpointError(protocol string) *hotPathEndpointError {
	if protocol == "anthropic" {
		return &hotPathEndpointError{Status: http.StatusBadGateway, Type: "api_error", Message: "workspace cleanup failed"}
	}
	return &hotPathEndpointError{Status: http.StatusBadGateway, Type: "run_error", Message: "workspace cleanup failed"}
}

// commitCleanupByLineage admits the exact cleanup continuation and removes the
// coordinator record in the same critical section. This is the terminal owner
// shared by cleanup-result and TTL races.
func (c *logicalRequestCoordinator) commitCleanupByLineage(
	ownerEdgeID, principalRef string,
	lineage logicalRequestContinuationLineage,
) (logicalRequestSnapshot, error) {
	c.mu.Lock()
	defer c.mu.Unlock()
	var target *logicalRequestRecord
	for _, record := range c.requests {
		if record.state == logicalRequestStateCleanup && record.cleanup && record.ownerEdgeID == ownerEdgeID && record.principalRef == principalRef &&
			record.lineage == lineage.Prefix && sameLogicalRequestResultIDs(record.expected, lineage.ResultIDs) {
			target = record
			break
		}
	}
	if target == nil {
		return logicalRequestSnapshot{}, errLogicalRequestNotFound
	}
	if err := validateLogicalRequestContinuationLineage(target.lineage, target.expectedIssuedCallHash, target.expected, lineage); err != nil {
		return logicalRequestSnapshot{}, err
	}
	snapshot := target.snapshot()
	delete(c.requests, target.id)
	return snapshot, nil
}

func (s *Server) writeHotPathTerminal(
	w http.ResponseWriter,
	r *http.Request,
	dispatch routeDispatch,
	protocol string,
	stream bool,
	requestID string,
	intent hotPathTerminalIntent,
) error {
	if intent.Error != nil {
		if protocol == "anthropic" {
			writeAnthropicError(w, intent.Error.Status, intent.Error.Type, intent.Error.Message)
		} else {
			writeError(w, intent.Error.Status, intent.Error.Type, intent.Error.Message)
		}
		return fmt.Errorf("%s", intent.Error.Message)
	}
	return s.writeHotPathStageResponse(w, r, dispatch, protocol, stream, requestID, intent.Output)
}

func hotPathLightEndpointError(protocol string, status int, message string) hotPathEndpointError {
	errorType := "run_error"
	if protocol == "anthropic" {
		errorType = "api_error"
	}
	return hotPathEndpointError{Status: status, Type: errorType, Message: message}
}

func (s *Server) retainHotPathPrimaryErrorForTTL(requestID string, primary hotPathEndpointError) *hotPathTerminalIntent {
	ownerEdgeID := s.edgeIDValue()
	if s.lightFlows != nil {
		s.lightFlows.abortDispatch(requestID, ownerEdgeID)
	}
	_ = s.requestCoordinator.disconnect(requestID, ownerEdgeID, "primary_error")
	return &hotPathTerminalIntent{Error: &primary}
}

func (s *Server) writeHotPathPrimaryError(
	w http.ResponseWriter,
	r *http.Request,
	dispatch routeDispatch,
	protocol string,
	stream bool,
	requestID string,
	primary hotPathEndpointError,
) error {
	ownerEdgeID := s.edgeIDValue()
	s.lightFlows.abortDispatch(requestID, ownerEdgeID)
	if err := r.Context().Err(); err != nil {
		s.disconnectHotPathRequest(requestID, ownerEdgeID)
		return err
	}

	cleanup, err := s.lightFlows.beginPrimaryErrorCleanup(r.Context(), requestID, ownerEdgeID, primary, s.requestCoordinator)
	if err == nil {
		return s.writeHotPathStageResponse(w, r, dispatch, protocol, stream, requestID, cleanup)
	}
	if contextErr := r.Context().Err(); contextErr != nil {
		s.disconnectHotPathRequest(requestID, ownerEdgeID)
		return contextErr
	}
	intent := s.retainHotPathPrimaryErrorForTTL(requestID, primary)
	return s.writeHotPathTerminal(w, r, dispatch, protocol, stream, requestID, *intent)
}

func (s *Server) disconnectHotPathRequest(requestID, ownerEdgeID string) {
	if requestID == "" {
		return
	}
	if s.lightFlows != nil {
		s.lightFlows.abortDispatch(requestID, ownerEdgeID)
	}
	_ = s.requestCoordinator.disconnect(requestID, ownerEdgeID, "cancelled")
}