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

package openai

import (
	"context"
	"errors"
	"sync"
	"time"

	edgeservice "iop/apps/edge/internal/service"
)

// SingleRequestExecutor is the concurrent, request-safe composite executor
// driving the Plan -> Work -> Review lifecycle behind the single-request interface.
type SingleRequestExecutor struct {
	provider *singleRequestProviderStage
	bridge   *singleRequestWorkToolBridge
	plan     *singleRequestPlanStage
	work     *singleRequestWorkStage
	review   *singleRequestReviewStage
}

// NewSingleRequestExecutor constructs a production composite single-request executor
// backed by private stage drivers and the correlated continuation bridge.
func NewSingleRequestExecutor(service edgeserviceRunner) *SingleRequestExecutor {
	provider := newSingleRequestProviderStage(service)
	bridge := newSingleRequestWorkToolBridge()
	return &SingleRequestExecutor{
		provider: provider,
		bridge:   bridge,
		plan:     newSingleRequestPlanStage(provider),
		work:     newSingleRequestWorkStage(provider, bridge),
		review:   newSingleRequestReviewStage(provider, bridge),
	}
}

// ExecuteSingleRequest executes Plan -> Work -> Review sequentially for one request
// against a single controller and immutable binding.
func (s *SingleRequestExecutor) ExecuteSingleRequest(ctx context.Context, req edgeservice.SingleRequestRequest, ctrl edgeservice.SingleRequestController) error {
	if s == nil || s.provider == nil || s.bridge == nil || s.plan == nil || s.work == nil || s.review == nil || ctrl == nil {
		return edgeservice.ErrSingleRequestExecutorUnavailable
	}
	if req.RequestID == "" || req.RequestID != ctrl.RequestID() {
		return edgeservice.ErrSingleRequestIdentityMismatch
	}
	binding := ctrl.Binding()
	if binding == nil || req.Binding == nil {
		return edgeservice.ErrSingleRequestInvalidBinding
	}
	if binding.Workspace == nil || binding.Workspace.NodeID == "" {
		return edgeservice.ErrSingleRequestInvalidBinding
	}

	defer s.bridge.clearRequest(req.RequestID)
	quality := newSingleRequestQualityGate()

	seqCtrl := &singleRequestSequenceController{
		SingleRequestController: ctrl,
	}

	nodeRef := binding.Workspace.NodeID

	// Stage 1: Plan
	planReq := singleRequestPlanStageRequest{
		RequestID:        req.RequestID,
		Task:             req.Prompt,
		StageBinding:     binding.Plan,
		Limits:           binding.Limits,
		NodeRef:          nodeRef,
		SessionID:        req.RequestID,
		UsageAttribution: singleRequestUsageAttribution(binding.Plan),
		Sequence:         1,
		Quality:          quality,
	}
	_, err := s.plan.run(ctx, planReq, seqCtrl)
	if err != nil {
		return submitSingleRequestClosedTerminal(ctx, req.RequestID, seqCtrl, err)
	}

	// Stage 2: Work
	workReq := singleRequestWorkStageRequest{
		RequestID:        req.RequestID,
		Task:             req.Prompt,
		StageBinding:     binding.Work,
		Limits:           binding.Limits,
		NodeRef:          nodeRef,
		SessionID:        req.RequestID,
		UsageAttribution: singleRequestUsageAttribution(binding.Work),
		Sequence:         1,
		Quality:          quality,
	}
	workResult, err := s.work.run(ctx, workReq, seqCtrl)
	if err != nil {
		return submitSingleRequestClosedTerminal(ctx, req.RequestID, seqCtrl, err)
	}

	// Stage 3: Review
	reviewReq := singleRequestReviewStageRequest{
		RequestID:        req.RequestID,
		Task:             req.Prompt,
		Work:             workResult,
		StageBinding:     binding.Review,
		Limits:           binding.Limits,
		NodeRef:          nodeRef,
		SessionID:        req.RequestID,
		UsageAttribution: singleRequestUsageAttribution(binding.Review),
		Sequence:         1,
		Quality:          quality,
	}
	_, err = s.review.run(ctx, reviewReq, seqCtrl)
	if err != nil {
		return submitSingleRequestClosedTerminal(ctx, req.RequestID, seqCtrl, err)
	}

	return nil
}

func submitSingleRequestClosedTerminal(ctx context.Context, requestID string, ctrl edgeservice.SingleRequestController, stageErr error) error {
	// The parent request context is owned by the service. Returning its error
	// lets that owner distinguish caller cancellation from wall-clock budget
	// exhaustion without racing a stage terminal envelope.
	if ctx != nil {
		if err := ctx.Err(); err != nil {
			return err
		}
		if deadline, ok := ctx.Deadline(); ok && !time.Now().Before(deadline) {
			return context.DeadlineExceeded
		}
	}
	disposition, ok := singleRequestTerminalDisposition(stageErr)
	if !ok {
		switch {
		case errors.Is(stageErr, context.DeadlineExceeded):
			disposition = edgeservice.SingleRequestTerminalDisposition{Kind: edgeservice.SingleRequestTerminalError, ErrorClass: edgeservice.SingleRequestTerminalErrorTimeout}
		default:
			disposition = edgeservice.SingleRequestTerminalDisposition{Kind: edgeservice.SingleRequestTerminalError, ErrorClass: edgeservice.SingleRequestTerminalErrorProvider}
		}
	}

	var envelope edgeservice.SingleRequestEnvelope
	switch disposition.Kind {
	case edgeservice.SingleRequestTerminalLength:
		envelope = edgeservice.SingleRequestEnvelope{
			RequestID: requestID,
			Stage:     edgeservice.SingleRequestStateFinalizing,
			Result: &edgeservice.SingleRequestResult{
				Terminal: disposition,
			},
		}
	case edgeservice.SingleRequestTerminalCancelled:
		envelope = edgeservice.SingleRequestEnvelope{RequestID: requestID, Stage: edgeservice.SingleRequestStateCancelled, Terminal: &disposition}
	default:
		if disposition.Kind != edgeservice.SingleRequestTerminalError || disposition.Validate() != nil {
			disposition = edgeservice.SingleRequestTerminalDisposition{Kind: edgeservice.SingleRequestTerminalError, ErrorClass: edgeservice.SingleRequestTerminalErrorProvider}
		}
		envelope = edgeservice.SingleRequestEnvelope{
			RequestID: requestID,
			Stage:     edgeservice.SingleRequestStateFailed,
			Terminal:  &disposition,
			Err:       edgeservice.ErrSingleRequestFailed,
		}
	}
	if err := ctrl.SubmitEnvelope(envelope); err != nil && !errors.Is(err, edgeservice.ErrSingleRequestTerminal) {
		return err
	}
	return nil
}

// ContinueInternalTool delegates correlated tool continuation results to the
// request-safe bridge.
func (s *SingleRequestExecutor) ContinueInternalTool(ctx context.Context, result edgeservice.InternalWorkspaceToolResult) error {
	if s == nil || s.bridge == nil {
		return edgeservice.ErrSingleRequestExecutorUnavailable
	}
	return s.bridge.ContinueInternalTool(ctx, result)
}

type singleRequestSequenceController struct {
	edgeservice.SingleRequestController
	mu       sync.Mutex
	sequence uint64
}

func (c *singleRequestSequenceController) SubmitEnvelope(env edgeservice.SingleRequestEnvelope) error {
	c.mu.Lock()
	c.sequence++
	env.Sequence = c.sequence
	c.mu.Unlock()
	return c.SingleRequestController.SubmitEnvelope(env)
}

func singleRequestUsageAttribution(binding edgeservice.SingleRequestStageBinding) string {
	if binding.Dispatch != nil {
		return binding.Dispatch.PrincipalRef
	}
	return ""
}