iop/apps/edge/internal/service/single_request_artifact.go

package service

import (
	"context"
	"errors"
	"time"

	iop "iop/proto/gen/iop"
)

var (
	ErrSingleRequestInternalArtifactUnavailable = errors.New("single-request internal artifact is unavailable")
	ErrSingleRequestInternalArtifactInvalid     = errors.New("single-request internal artifact request is invalid")
	ErrSingleRequestInternalArtifactBudget      = errors.New("single-request internal artifact budget exceeded")
	ErrSingleRequestInternalArtifactFailed      = errors.New("single-request internal artifact operation failed")
	ErrSingleRequestWorkspaceOpen               = errors.New("single-request workspace open failed")
)

// SingleRequestArtifactKind is the coordinator-facing closed artifact set. It
// deliberately carries no filename or relative path.
type SingleRequestArtifactKind string

const (
	SingleRequestArtifactPlan   SingleRequestArtifactKind = "plan"
	SingleRequestArtifactReview SingleRequestArtifactKind = "review"
)

type singleRequestWorkspaceArtifactRuntime interface {
	workspaceArtifact(context.Context, *SingleRequestWorkspaceBinding, *iop.WorkspaceArtifactRequest, int) (*iop.WorkspaceArtifactResponse, error)
}

type singleRequestArtifactOperation struct {
	ctx         context.Context
	cancel      context.CancelFunc
	runtime     singleRequestWorkspaceArtifactRuntime
	openRuntime singleRequestWorkspaceToolRuntime
	binding     *SingleRequestWorkspaceBinding
	kind        iop.WorkspaceArtifactKind
	maximum     int
	complete    func()
}

func (h *singleRequestHandle) ReadInternalArtifact(ctx context.Context, kind SingleRequestArtifactKind) ([]byte, error) {
	operation, err := h.prepareSingleRequestArtifact(ctx, kind, 0)
	if err != nil {
		h.failSingleRequestArtifact(ctx, err)
		return nil, err
	}
	defer operation.cancel()
	defer operation.complete()

	if err := operation.ctx.Err(); err != nil {
		return nil, h.finishSingleRequestArtifactFailure(operation.ctx, err)
	}
	if err := h.ensureSingleRequestWorkspaceOpen(operation.ctx, operation.openRuntime, operation.binding); err != nil {
		return nil, h.finishSingleRequestArtifactFailure(operation.ctx, err)
	}
	if err := operation.ctx.Err(); err != nil {
		return nil, h.finishSingleRequestArtifactFailure(operation.ctx, err)
	}
	response, err := operation.runtime.workspaceArtifact(operation.ctx, operation.binding, &iop.WorkspaceArtifactRequest{
		RequestId: h.req.RequestID,
		Kind:      operation.kind,
		Operation: iop.WorkspaceArtifactOperation_WORKSPACE_ARTIFACT_OPERATION_READ,
	}, operation.maximum)
	if err != nil || response.GetStatus() != iop.WorkspaceStatus_WORKSPACE_STATUS_SUCCESS {
		return nil, h.finishSingleRequestArtifactFailure(operation.ctx, err)
	}
	if err := operation.ctx.Err(); err != nil {
		return nil, h.finishSingleRequestArtifactFailure(operation.ctx, err)
	}
	return append([]byte(nil), response.GetContent()...), nil
}

func (h *singleRequestHandle) WriteInternalArtifact(ctx context.Context, kind SingleRequestArtifactKind, content []byte) error {
	operation, err := h.prepareSingleRequestArtifact(ctx, kind, len(content))
	if err != nil {
		h.failSingleRequestArtifact(ctx, err)
		return err
	}
	defer operation.cancel()
	defer operation.complete()

	if err := operation.ctx.Err(); err != nil {
		return h.finishSingleRequestArtifactFailure(operation.ctx, err)
	}
	if err := h.ensureSingleRequestWorkspaceOpen(operation.ctx, operation.openRuntime, operation.binding); err != nil {
		return h.finishSingleRequestArtifactFailure(operation.ctx, err)
	}
	if err := operation.ctx.Err(); err != nil {
		return h.finishSingleRequestArtifactFailure(operation.ctx, err)
	}
	response, err := operation.runtime.workspaceArtifact(operation.ctx, operation.binding, &iop.WorkspaceArtifactRequest{
		RequestId: h.req.RequestID,
		Kind:      operation.kind,
		Operation: iop.WorkspaceArtifactOperation_WORKSPACE_ARTIFACT_OPERATION_WRITE,
		Content:   append([]byte(nil), content...),
	}, operation.maximum)
	if err != nil || response.GetStatus() != iop.WorkspaceStatus_WORKSPACE_STATUS_SUCCESS {
		return h.finishSingleRequestArtifactFailure(operation.ctx, err)
	}
	if err := operation.ctx.Err(); err != nil {
		return h.finishSingleRequestArtifactFailure(operation.ctx, err)
	}
	return nil
}

func (h *singleRequestHandle) prepareSingleRequestArtifact(ctx context.Context, kind SingleRequestArtifactKind, contentBytes int) (*singleRequestArtifactOperation, error) {
	if ctx == nil {
		return nil, ErrSingleRequestInternalArtifactInvalid
	}
	if err := ctx.Err(); err != nil {
		return nil, err
	}
	protoKind, ok := singleRequestArtifactProtoKind(kind)
	if !ok {
		return nil, ErrSingleRequestInternalArtifactInvalid
	}

	h.mu.Lock()
	if isTerminalState(h.state) || h.state == SingleRequestStateFinalizing {
		h.mu.Unlock()
		return nil, ErrSingleRequestTerminal
	}
	if h.state == SingleRequestStateInternalTool || h.activeStageIDLocked() == "" || h.toolLoop.stageDeadline.IsZero() {
		h.mu.Unlock()
		return nil, ErrSingleRequestInternalArtifactInvalid
	}
	openRuntime := h.toolLoop.runtime
	runtime, ok := openRuntime.(singleRequestWorkspaceArtifactRuntime)
	binding := h.binding.Workspace.Clone()
	deadline := h.toolLoop.stageDeadline
	if !ok || runtime == nil || binding == nil {
		h.mu.Unlock()
		return nil, ErrSingleRequestInternalArtifactUnavailable
	}
	maximum := binding.Limits.MaxOutputBytes
	if maximum < 1 || contentBytes < 0 || contentBytes > maximum {
		h.mu.Unlock()
		return nil, ErrSingleRequestInternalArtifactBudget
	}
	if !time.Now().Before(deadline) {
		h.mu.Unlock()
		return nil, ErrSingleRequestInternalArtifactBudget
	}
	h.toolWg.Add(1)
	h.toolWork++
	h.mu.Unlock()

	operationCtx, cancel := singleRequestArtifactContext(h.execCtx, ctx, deadline)
	return &singleRequestArtifactOperation{
		ctx: operationCtx, cancel: cancel, runtime: runtime, openRuntime: openRuntime, binding: binding,
		kind: protoKind, maximum: maximum,
		complete: func() {
			h.mu.Lock()
			h.toolWork--
			h.mu.Unlock()
			h.toolWg.Done()
		},
	}, nil
}

func singleRequestArtifactProtoKind(kind SingleRequestArtifactKind) (iop.WorkspaceArtifactKind, bool) {
	switch kind {
	case SingleRequestArtifactPlan:
		return iop.WorkspaceArtifactKind_WORKSPACE_ARTIFACT_KIND_PLAN, true
	case SingleRequestArtifactReview:
		return iop.WorkspaceArtifactKind_WORKSPACE_ARTIFACT_KIND_REVIEW, true
	default:
		return iop.WorkspaceArtifactKind_WORKSPACE_ARTIFACT_KIND_UNSPECIFIED, false
	}
}

func singleRequestArtifactContext(execution, caller context.Context, stageDeadline time.Time) (context.Context, context.CancelFunc) {
	deadline := stageDeadline
	if callerDeadline, ok := caller.Deadline(); ok && callerDeadline.Before(deadline) {
		deadline = callerDeadline
	}
	ctx, cancel := context.WithDeadline(execution, deadline)
	stop := context.AfterFunc(caller, cancel)
	return ctx, func() {
		stop()
		cancel()
	}
}

func (h *singleRequestHandle) finishSingleRequestArtifactFailure(ctx context.Context, _ error) error {
	err := ErrSingleRequestInternalArtifactFailed
	if ctx != nil && ctx.Err() != nil {
		err = ctx.Err()
	}
	h.failSingleRequestArtifact(ctx, err)
	return err
}

func (h *singleRequestHandle) failSingleRequestArtifact(ctx context.Context, err error) {
	fallback := singleRequestErrorClassFromErr(err)
	if errors.Is(err, context.DeadlineExceeded) || errors.Is(err, ErrSingleRequestInternalArtifactBudget) {
		fallback = singleRequestErrorClassTimeout
	}
	outcome, errorClass := h.classifyChildOperationContext(ctx, fallback)
	h.mu.Lock()
	defer h.mu.Unlock()
	if isTerminalState(h.state) || h.state == SingleRequestStateFinalizing {
		return
	}
	if outcome == singleRequestOutcomeCancel {
		h.cancelLocked()
		return
	}
	if errorClass == singleRequestErrorClassInternalToolBudget {
		h.failLockedWithErrorClass(ErrSingleRequestInternalToolBudget, singleRequestErrorClassInternalToolBudget)
		return
	}
	if errorClass == singleRequestErrorClassTimeout {
		h.failLockedWithErrorClass(ErrSingleRequestInternalArtifactBudget, singleRequestErrorClassTimeout)
		return
	}
	h.failLockedWithErrorClass(err, errorClass)
}

// ensureSingleRequestWorkspaceOpen serializes the one permitted open attempt
// across model-tool and internal-artifact callers. A failed attempt is cached so
// no racing caller can silently send a second open on the same request.
func (h *singleRequestHandle) ensureSingleRequestWorkspaceOpen(ctx context.Context, runtime singleRequestWorkspaceToolRuntime, binding *SingleRequestWorkspaceBinding) error {
	for {
		if err := ctx.Err(); err != nil {
			return err
		}
		h.mu.Lock()
		if h.toolLoop.opened {
			h.mu.Unlock()
			return nil
		}
		if h.toolLoop.opening {
			done := h.toolLoop.openDone
			h.mu.Unlock()
			select {
			case <-done:
				continue
			case <-ctx.Done():
				return ctx.Err()
			}
		}
		if h.toolLoop.openAttempted {
			err := h.toolLoop.openErr
			if err == nil {
				err = ErrSingleRequestWorkspaceOpen
			}
			h.mu.Unlock()
			return err
		}
		if runtime == nil || binding == nil {
			h.mu.Unlock()
			return ErrSingleRequestWorkspaceOpen
		}
		h.toolLoop.openAttempted = true
		h.toolLoop.opening = true
		h.toolLoop.openDone = make(chan struct{})
		done := h.toolLoop.openDone
		h.mu.Unlock()

		response, err := runtime.workspaceOpen(ctx, binding, &iop.WorkspaceOpenRequest{
			RequestId: h.req.RequestID, WorkspaceRef: binding.Ref,
			TimeoutMs: singleRequestContextRemainingMilliseconds(ctx),
		})
		accepted := err == nil && response.GetStatus() == iop.WorkspaceStatus_WORKSPACE_STATUS_SUCCESS
		h.mu.Lock()
		if accepted {
			h.toolLoop.opened = true
			h.toolLoop.openErr = nil
		} else {
			h.toolLoop.openErr = ErrSingleRequestWorkspaceOpen
		}
		h.toolLoop.opening = false
		close(done)
		h.mu.Unlock()
		if !accepted {
			if ctx.Err() != nil {
				return ctx.Err()
			}
			return ErrSingleRequestWorkspaceOpen
		}
		if ctx.Err() != nil {
			return ctx.Err()
		}
		return nil
	}
}

func singleRequestContextRemainingMilliseconds(ctx context.Context) int64 {
	deadline, ok := ctx.Deadline()
	if !ok {
		return 0
	}
	return internalToolRemainingMilliseconds(deadline)
}