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

package service

import (
	"context"
	"errors"
	"slices"
	"strings"
	"time"

	iop "iop/proto/gen/iop"
)

type singleRequestWorkspaceToolRuntime interface {
	workspaceOpen(context.Context, *SingleRequestWorkspaceBinding, *iop.WorkspaceOpenRequest) (*iop.WorkspaceOpenResponse, error)
	workspaceTool(context.Context, *SingleRequestWorkspaceBinding, *iop.WorkspaceToolRequest) (*iop.WorkspaceToolResponse, error)
}

// SingleRequestWorkspaceLifecycle is optional for executors that never open a
// workspace. Once a workspace is open, terminal commit fails closed unless the
// lifecycle can complete one typed cleanup for the immutable request.
type SingleRequestWorkspaceLifecycle interface {
	CleanupWorkspace(context.Context, *SingleRequestWorkspaceBinding, string) error
}

type singleRequestToolUsage struct {
	iterations  int
	outputBytes int
}

type singleRequestToolLoopState struct {
	continuation       SingleRequestToolContinuation
	runtime            singleRequestWorkspaceToolRuntime
	lifecycle          SingleRequestWorkspaceLifecycle
	opened             bool
	seenCallIDs        map[string]struct{}
	usage              map[string]singleRequestToolUsage
	pendingCallID      string
	pendingResultReady bool
	stageID            string
	stageDeadline      time.Time
	stageEpoch         uint64
	stageTimer         *time.Timer
}

type singleRequestPendingTool struct {
	request  *iop.WorkspaceToolRequest
	stageID  string
	deadline time.Time
}

func (h *singleRequestHandle) prepareInternalWorkspaceToolLocked(call *InternalWorkspaceToolCall) (*singleRequestPendingTool, error, singleRequestErrorClass) {
	if h.toolLoop.continuation == nil || h.toolLoop.runtime == nil {
		return nil, ErrSingleRequestInternalToolUnavailable, ""
	}
	if h.toolLoop.pendingCallID != "" {
		return nil, ErrSingleRequestInternalToolInvalidCall, ""
	}
	cloned := call.Clone()
	request, err := decodeInternalWorkspaceToolCall(cloned)
	if err != nil {
		return nil, err, ""
	}
	expectedStageID := canonicalSingleRequestStageID(h.state)
	if expectedStageID == "" || request.GetRequestId() != h.req.RequestID || request.GetStageId() != expectedStageID {
		return nil, ErrSingleRequestIdentityMismatch, ""
	}
	if h.binding == nil || h.binding.Workspace == nil || !h.internalWorkspaceToolCapabilityAllowed(request) {
		return nil, ErrSingleRequestInternalToolDenied, ""
	}
	if _, duplicate := h.toolLoop.seenCallIDs[request.GetToolCallId()]; duplicate {
		return nil, ErrSingleRequestInternalToolInvalidCall, ""
	}
	usage := h.toolLoop.usage[expectedStageID]
	if usage.iterations >= h.binding.Limits.MaxToolIterations || h.toolLoop.stageDeadline.IsZero() {
		return nil, ErrSingleRequestInternalToolBudget, ""
	}
	if !time.Now().Before(h.toolLoop.stageDeadline) {
		return nil, ErrSingleRequestInternalToolBudget, singleRequestErrorClassTimeout
	}

	usage.iterations++
	h.toolLoop.usage[expectedStageID] = usage
	h.toolLoop.seenCallIDs[request.GetToolCallId()] = struct{}{}
	h.toolLoop.pendingCallID = request.GetToolCallId()
	h.toolLoop.pendingResultReady = false
	return &singleRequestPendingTool{
		request:  request,
		stageID:  expectedStageID,
		deadline: h.toolLoop.stageDeadline,
	}, nil, ""
}

func (h *singleRequestHandle) internalWorkspaceToolCapabilityAllowed(request *iop.WorkspaceToolRequest) bool {
	workspace := h.binding.Workspace
	operationID := ""
	switch request.GetOperation() {
	case iop.WorkspaceOperation_WORKSPACE_OPERATION_READ:
		operationID = "read"
	case iop.WorkspaceOperation_WORKSPACE_OPERATION_LIST:
		operationID = "list"
	case iop.WorkspaceOperation_WORKSPACE_OPERATION_WRITE:
		operationID = "write"
		if write := request.GetWrite(); write == nil || len(write.GetContent()) > workspace.Limits.MaxWriteBytes {
			return false
		}
	case iop.WorkspaceOperation_WORKSPACE_OPERATION_DELETE:
		operationID = "delete"
	case iop.WorkspaceOperation_WORKSPACE_OPERATION_COMMAND:
		operationID = "command"
		if _, allowed := slices.BinarySearch(workspace.CommandIDs, request.GetCommandId()); !allowed {
			return false
		}
		environmentBytes := 0
		for name, value := range request.GetEnvironment() {
			if _, allowed := slices.BinarySearch(workspace.EnvironmentNames, name); !allowed || strings.ContainsRune(value, 0) {
				return false
			}
			environmentBytes += len(name) + len(value)
			if environmentBytes > h.binding.Limits.MaxOutputBytes {
				return false
			}
		}
	default:
		return false
	}
	_, allowed := slices.BinarySearch(workspace.OperationIDs, operationID)
	return allowed
}

func (h *singleRequestHandle) executeInternalWorkspaceTool(pending *singleRequestPendingTool) {
	outcome := singleRequestOutcomeSuccess
	errorClass := singleRequestErrorClass("")
	defer func() {
		h.mu.Lock()
		h.timing.onToolExit(outcome, errorClass)
		h.mu.Unlock()
	}()
	if pending == nil || pending.request == nil {
		outcome, errorClass = singleRequestOutcomeError, singleRequestErrorClassValidation
		h.failInternalWorkspaceTool(ErrSingleRequestInternalToolInvalidCall)
		return
	}
	ctx, cancel := context.WithDeadline(h.execCtx, pending.deadline)
	defer cancel()

	h.mu.Lock()
	needOpen := !h.toolLoop.opened
	runtime := h.toolLoop.runtime
	continuation := h.toolLoop.continuation
	binding := h.binding.Workspace.Clone()
	h.mu.Unlock()
	if runtime == nil || continuation == nil || binding == nil {
		outcome, errorClass = singleRequestOutcomeError, singleRequestErrorClassProvider
		h.failInternalWorkspaceTool(ErrSingleRequestInternalToolUnavailable)
		return
	}

	if needOpen {
		openResponse, err := runtime.workspaceOpen(ctx, binding, &iop.WorkspaceOpenRequest{
			RequestId:    h.req.RequestID,
			WorkspaceRef: binding.Ref,
			TimeoutMs:    internalToolRemainingMilliseconds(pending.deadline),
		})
		if err != nil || openResponse.GetStatus() != iop.WorkspaceStatus_WORKSPACE_STATUS_SUCCESS {
			outcome, errorClass = singleRequestToolOutcome(ctx)
			h.failInternalWorkspaceToolOutcome(ctx, err)
			return
		}
		h.mu.Lock()
		if h.toolLoop.pendingCallID == pending.request.GetToolCallId() {
			h.toolLoop.opened = true
		}
		terminal := isTerminalState(h.state)
		h.mu.Unlock()
		if terminal {
			return
		}
	}

	pending.request.TimeoutMs = 0
	if pending.request.GetOperation() == iop.WorkspaceOperation_WORKSPACE_OPERATION_COMMAND {
		pending.request.TimeoutMs = internalToolCommandTimeoutMilliseconds(pending.deadline, binding.Limits.MaxCommandTimeoutMS)
	}
	response, err := runtime.workspaceTool(ctx, binding, pending.request)
	if err != nil {
		outcome, errorClass = singleRequestToolOutcome(ctx)
		h.failInternalWorkspaceToolOutcome(ctx, err)
		return
	}
	if response.GetStatus() != iop.WorkspaceStatus_WORKSPACE_STATUS_SUCCESS &&
		response.GetErrorCode() != iop.WorkspaceErrorCode_WORKSPACE_ERROR_CODE_NOT_FOUND {
		outcome, errorClass = singleRequestOutcomeError, singleRequestErrorClassInternalToolFailed
		h.failInternalWorkspaceTool(ErrSingleRequestInternalToolFailed)
		return
	}
	result := internalWorkspaceToolResult(response)

	h.mu.Lock()
	if isTerminalState(h.state) {
		outcome, errorClass = singleRequestToolOutcome(ctx)
		h.mu.Unlock()
		return
	}
	if h.state != SingleRequestStateInternalTool || h.savedStage == "" || h.toolLoop.pendingCallID != result.ToolCallID ||
		canonicalSingleRequestStageID(h.savedStage) != pending.stageID || result.RequestID != h.req.RequestID || result.StageID != pending.stageID {
		outcome, errorClass = singleRequestOutcomeError, singleRequestErrorClassValidation
		h.failLocked(ErrSingleRequestIdentityMismatch)
		h.mu.Unlock()
		return
	}
	usage := h.toolLoop.usage[pending.stageID]
	outputBytes := internalWorkspaceToolOutputBytes(result)
	if outputBytes < 0 || outputBytes > h.binding.Limits.MaxOutputBytes-usage.outputBytes {
		outcome, errorClass = singleRequestOutcomeError, singleRequestErrorClassInternalToolBudget
		h.failLocked(ErrSingleRequestInternalToolBudget)
		h.mu.Unlock()
		return
	}
	usage.outputBytes += outputBytes
	h.toolLoop.usage[pending.stageID] = usage
	h.toolLoop.pendingResultReady = true
	h.mu.Unlock()

	if err := continuation.ContinueInternalTool(ctx, result.Clone()); err != nil {
		outcome, errorClass = singleRequestOutcomeError, singleRequestErrorClassInternalToolFailed
		h.failInternalWorkspaceTool(ErrSingleRequestInternalToolFailed)
	}
}

func singleRequestToolOutcome(ctx context.Context) (singleRequestOutcome, singleRequestErrorClass) {
	if deadline, ok := ctx.Deadline(); ok && !time.Now().Before(deadline) {
		return singleRequestOutcomeError, singleRequestErrorClassTimeout
	}
	if errors.Is(ctx.Err(), context.Canceled) {
		return singleRequestOutcomeCancel, singleRequestErrorClassCancel
	}
	if errors.Is(ctx.Err(), context.DeadlineExceeded) {
		return singleRequestOutcomeError, singleRequestErrorClassTimeout
	}
	return singleRequestOutcomeError, singleRequestErrorClassInternalToolFailed
}

// CleanupWorkspace maps the private wire terminal to one safe coordinator
// outcome. Node error text and filesystem details never enter coordinator state.
func (s *Service) CleanupWorkspace(ctx context.Context, binding *SingleRequestWorkspaceBinding, requestID string) error {
	response, err := s.workspaceCleanup(ctx, binding, &iop.WorkspaceCleanupRequest{RequestId: requestID})
	if err != nil || response.GetStatus() != iop.WorkspaceStatus_WORKSPACE_STATUS_SUCCESS {
		return ErrSingleRequestWorkspaceCleanup
	}
	return nil
}

func (h *singleRequestHandle) failInternalWorkspaceToolOutcome(ctx context.Context, err error) {
	deadline, hasDeadline := ctx.Deadline()
	if errors.Is(ctx.Err(), context.DeadlineExceeded) || hasDeadline && !time.Now().Before(deadline) {
		h.failInternalWorkspaceToolWithErrorClass(ErrSingleRequestInternalToolBudget, singleRequestErrorClassTimeout)
		return
	}
	if errors.Is(ctx.Err(), context.Canceled) {
		h.mu.Lock()
		if !isTerminalState(h.state) {
			h.cancelLocked()
		}
		h.mu.Unlock()
		return
	}
	_ = err
	h.failInternalWorkspaceTool(ErrSingleRequestInternalToolFailed)
}

func (h *singleRequestHandle) failInternalWorkspaceTool(err error) {
	h.failInternalWorkspaceToolWithErrorClass(err, "")
}

func (h *singleRequestHandle) failInternalWorkspaceToolWithErrorClass(err error, errorClass singleRequestErrorClass) {
	h.mu.Lock()
	if !isTerminalState(h.state) {
		h.failLockedWithErrorClass(err, errorClass)
	}
	h.mu.Unlock()
}

func (h *singleRequestHandle) activeStageIDLocked() string {
	if h.state == SingleRequestStateInternalTool {
		return canonicalSingleRequestStageID(h.savedStage)
	}
	return canonicalSingleRequestStageID(h.state)
}

func canonicalSingleRequestStageID(state SingleRequestState) string {
	switch state {
	case SingleRequestStatePlanning:
		return "plan"
	case SingleRequestStateWorking:
		return "work"
	case SingleRequestStateReviewing, SingleRequestStateRepairing:
		return "review"
	default:
		return ""
	}
}

func (h *singleRequestHandle) updateStageBudgetLocked(previousStageID string) {
	nextStageID := h.activeStageIDLocked()
	if previousStageID == nextStageID {
		return
	}
	h.stopStageBudgetLocked()
	h.toolLoop.stageID = nextStageID
	if nextStageID == "" {
		h.toolLoop.stageDeadline = time.Time{}
		return
	}
	h.toolLoop.stageDeadline = time.Now().Add(time.Duration(h.binding.Limits.StageTimeoutMS) * time.Millisecond)
	h.toolLoop.stageEpoch++
	epoch := h.toolLoop.stageEpoch
	deadline := h.toolLoop.stageDeadline
	h.toolLoop.stageTimer = time.AfterFunc(time.Until(deadline), func() {
		h.mu.Lock()
		defer h.mu.Unlock()
		if isTerminalState(h.state) || h.toolLoop.stageEpoch != epoch || h.activeStageIDLocked() != nextStageID {
			return
		}
		h.failLockedWithErrorClass(ErrSingleRequestInternalToolBudget, singleRequestErrorClassTimeout)
	})
}

func (h *singleRequestHandle) stopStageBudgetLocked() {
	if h.toolLoop.stageTimer != nil {
		h.toolLoop.stageTimer.Stop()
		h.toolLoop.stageTimer = nil
	}
}

func internalToolRemainingMilliseconds(deadline time.Time) int64 {
	remaining := time.Until(deadline)
	if remaining <= time.Millisecond {
		return 1
	}
	return int64((remaining + time.Millisecond - 1) / time.Millisecond)
}

func internalToolCommandTimeoutMilliseconds(deadline time.Time, workspaceMaximum int) int64 {
	remaining := internalToolRemainingMilliseconds(deadline)
	if workspaceMaximum > 0 && int64(workspaceMaximum) < remaining {
		return int64(workspaceMaximum)
	}
	return remaining
}