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

package openai

import (
	"bytes"
	"context"
	"fmt"
	"io"
	"net/http"
	"strings"
	"sync"
	"time"

	edgeservice "iop/apps/edge/internal/service"
	"iop/packages/go/config"
	iop "iop/proto/gen/iop"
)

// providerFakeRunService extends the normalized fakeRunService with the
// provider tunnel/pool dispatch state used by provider scenario tests.
type providerFakeRunService struct {
	fakeRunService

	// tunnelProviderURL, when set, makes SubmitProviderTunnel relay the tunnel
	// request to this httptest provider fixture and stream the raw provider
	// response back as ordered frames (emulating the Node relay).
	tunnelProviderURL string
	// tunnelServedTarget emulates provider-pool admission: BuildBody is called
	// with this target when non-empty.
	tunnelServedTarget string
	// tunnelFrames, when set, is handed to the tunnel handle as-is instead of
	// contacting tunnelProviderURL.
	tunnelFrames chan *iop.ProviderTunnelFrame
	tunnelErr    error
	tunnelReqs   []edgeservice.SubmitProviderTunnelRequest
	tunnelBodies [][]byte
	// tunnelWaitTimeout overrides the tunnel handle wait timeout for
	// timeout-path fixtures; zero keeps the 5s default.
	tunnelWaitTimeout time.Duration

	// poolDispatchPath records which execution path was returned by
	// SubmitProviderPool: "provider_tunnel" or "normalized".
	poolDispatchPath string
	// poolRunFrames is the frame source used by the normalized path when
	// SubmitProviderPool is in use.
	poolRunFrames chan *iop.RunEvent
	// lastTunnelHandle stores the tunnel handle from the last SubmitProviderPool
	// tunnel path call, so tests can check headers set by the handler.
	lastTunnelHandle *fakeTunnelHandle
	// poolSubmitResults provides custom results for successive SubmitProviderPool
	// calls. When non-empty, results are consumed in order; when empty, the
	// poolDispatchPath / tunnelErr fields are used instead.
	poolSubmitResults []edgeservice.ProviderPoolDispatchResult
	// poolSubmitCount tracks how many times SubmitProviderPool was called
	// (including from the chat handler). Used by tests to assert dispatch
	// surface fidelity.
	poolSubmitCount int
	// poolLastRun records the last SubmitRunRequest passed to SubmitProviderPool
	// (pre-PrepareRun). Tests assert ModelGroupKey is non-empty to catch the
	// "empty Run" regression.
	poolLastRun edgeservice.SubmitRunRequest
	// poolPrepareRunCalled tracks whether the optional PrepareRun hook was
	// invoked by SubmitProviderPool on the normalized path. Only meaningful
	// when poolDispatchPath == "normalized".
	poolPrepareRunCalled bool
	// poolSelectedCandidate is passed to PrepareProtocolTunnel. Zero value
	// emulates a legacy candidate without a concrete profile.
	poolSelectedCandidate edgeservice.ProviderPoolCandidate
}

type fakeTunnelHandle struct {
	dispatch    edgeservice.RunDispatch
	headers     map[string]string
	frames      chan *iop.ProviderTunnelFrame
	waitTimeout time.Duration
}

func (h *fakeTunnelHandle) Headers() map[string]string        { return h.headers }
func (h *fakeTunnelHandle) SetHeaders(hdrs map[string]string) { h.headers = hdrs }
func (h *fakeTunnelHandle) Dispatch() edgeservice.RunDispatch {
	disp := h.dispatch
	if strings.TrimSpace(disp.ProviderID) == "" {
		disp.ProviderID = "test-provider"
	}
	if strings.TrimSpace(disp.NodeID) == "" {
		disp.NodeID = "test-node"
	}
	return disp
}
func (h *fakeTunnelHandle) Stream() edgeservice.ProviderTunnelStream {
	return edgeservice.ProviderTunnelStream{Frames: h.frames}
}
func (h *fakeTunnelHandle) Close() {}
func (h *fakeTunnelHandle) WaitTimeout() time.Duration {
	if h.waitTimeout > 0 {
		return h.waitTimeout
	}
	return 5 * time.Second
}

// fakeRunResult implements edgeservice.RunResult for tests.
type fakeRunResult struct {
	dispatch edgeservice.RunDispatch
	events   chan *iop.RunEvent
}

func (r *fakeRunResult) Dispatch() edgeservice.RunDispatch {
	disp := r.dispatch
	if strings.TrimSpace(disp.ProviderID) == "" {
		disp.ProviderID = "test-provider"
	}
	if strings.TrimSpace(disp.NodeID) == "" {
		disp.NodeID = "test-node"
	}
	return disp
}
func (r *fakeRunResult) Stream() edgeservice.RunStream {
	return edgeservice.RunStream{Events: r.events}
}
func (r *fakeRunResult) Close() {}
func (r *fakeRunResult) WaitTimeout() time.Duration {
	return 5 * time.Second
}

func (s *providerFakeRunService) SubmitProviderTunnel(_ context.Context, req edgeservice.SubmitProviderTunnelRequest) (edgeservice.ProviderTunnelResult, error) {
	s.submitMu.Lock()
	s.tunnelReqs = append(s.tunnelReqs, req)
	s.submitMu.Unlock()
	if s.tunnelErr != nil {
		return nil, s.tunnelErr
	}
	target := req.Target
	if s.tunnelServedTarget != "" {
		target = s.tunnelServedTarget
	}
	body := req.Body
	if req.BuildBody != nil {
		built, err := req.BuildBody(target)
		if err != nil {
			return nil, err
		}
		body = built
	}
	s.submitMu.Lock()
	s.tunnelBodies = append(s.tunnelBodies, body)
	s.submitMu.Unlock()

	frames := s.tunnelFrames
	if frames == nil && s.tunnelProviderURL != "" {
		relayed, err := relayProviderFixtureFrames(s.tunnelProviderURL, req.Method, req.Path, body)
		if err != nil {
			return nil, err
		}
		frames = relayed
	}
	providerID := strings.TrimSpace(req.ProviderID)
	if providerID == "" {
		providerID = "test-provider"
	}
	nodeID := strings.TrimSpace(req.NodeRef)
	if nodeID == "" {
		nodeID = "node-1"
	}
	attribution := strings.TrimSpace(req.UsageAttribution)
	if attribution == "" {
		attribution = config.UsageAttributionProvider
	}
	return &fakeTunnelHandle{
		dispatch: edgeservice.RunDispatch{
			RunID:            "run-tunnel",
			NodeID:           nodeID,
			ProviderID:       providerID,
			UsageAttribution: attribution,
			ModelGroupKey:    req.ModelGroupKey,
			Adapter:          req.Adapter,
			Target:           target,
			SessionID:        req.SessionID,
			TimeoutSec:       5,
		},
		frames:      frames,
		waitTimeout: s.tunnelWaitTimeout,
	}, nil
}

func checkFakeCandidateAdmission(candidate edgeservice.ProviderPoolCandidate, req edgeservice.ProviderPoolDispatchRequest) error {
	operation := strings.TrimSpace(req.Tunnel.Operation)
	if operation != "" && candidate.ProtocolProfile != nil {
		if _, supported := candidate.ProtocolProfile.Operations[operation]; !supported {
			return &edgeservice.ProviderPoolOperationUnsupportedError{Operation: operation}
		}
	}
	if req.AcceptCandidate != nil && !req.AcceptCandidate(candidate) {
		return edgeservice.ErrProviderPoolCandidateRejected
	}
	return nil
}

// SubmitProviderPool is a minimal test double for the provider-pool one-shot
// dispatch. It records the selected path (tunnel vs normalized) and returns a
// stub result so the OpenAI handler can exercise the selection-first flow.
// It also tracks whether the optional PrepareRun hook was invoked, so tests
// can assert the handler only calls PrepareRun on the normalized path.
func (s *providerFakeRunService) SubmitProviderPool(_ context.Context, req edgeservice.ProviderPoolDispatchRequest) (*edgeservice.ProviderPoolDispatchResult, error) {
	s.submitMu.Lock()
	s.poolSubmitCount++
	s.poolLastRun = req.Run
	s.poolPrepareRunCalled = false
	s.submitMu.Unlock()

	// Custom results for successive calls (e.g. tool-validation retry tests).
	s.submitMu.Lock()
	if len(s.poolSubmitResults) > 0 {
		result := s.poolSubmitResults[0]
		s.poolSubmitResults = s.poolSubmitResults[1:]
		s.submitMu.Unlock()
		if err := checkFakeCandidateAdmission(s.poolSelectedCandidate, req); err != nil {
			return nil, err
		}
		// Apply pre-dispatch tunnel preparation for tunnel path when custom
		// results are used.
		if result.Path == edgeservice.ProviderPoolPathTunnel && (req.PrepareProtocolTunnel != nil || req.PrepareTunnel != nil) {
			tunnelReqPrepared := req.Tunnel
			var prepErr error
			if req.PrepareProtocolTunnel != nil {
				tunnelReqPrepared, prepErr = req.PrepareProtocolTunnel(tunnelReqPrepared, s.poolSelectedCandidate)
			} else {
				tunnelReqPrepared, prepErr = req.PrepareTunnel(tunnelReqPrepared)
			}
			if prepErr != nil {
				return nil, prepErr
			}
			req.Tunnel = tunnelReqPrepared
		}
		return s.buildPoolResult(req, result)
	}
	err := s.tunnelErr
	path := s.poolDispatchPath
	s.submitMu.Unlock()
	if err := checkFakeCandidateAdmission(s.poolSelectedCandidate, req); err != nil {
		return nil, err
	}

	// Apply pre-dispatch tunnel preparation (mirrors real SubmitProviderPool
	// behavior: PrepareTunnel runs BEFORE buildProviderTunnelRequest and the
	// Node Send step). On failure, no tunnel request is recorded or sent.
	tunnelReq := req.Tunnel
	if path == string(edgeservice.ProviderPoolPathTunnel) || path == "" {
		if selectedProvider := strings.TrimSpace(s.poolSelectedCandidate.ProviderID); selectedProvider != "" {
			tunnelReq.ProviderID = selectedProvider
		}
		if selectedTarget := strings.TrimSpace(s.poolSelectedCandidate.ActualModel); selectedTarget != "" {
			tunnelReq.Target = selectedTarget
		}
		if req.PrepareProtocolTunnel != nil {
			tunnelReqPrepared, prepErr := req.PrepareProtocolTunnel(tunnelReq, s.poolSelectedCandidate)
			if prepErr != nil {
				return nil, prepErr
			}
			tunnelReq = tunnelReqPrepared
		} else if req.PrepareTunnel != nil {
			tunnelReqPrepared, prepErr := req.PrepareTunnel(tunnelReq)
			if prepErr != nil {
				return nil, prepErr
			}
			tunnelReq = tunnelReqPrepared
		}
		req.Tunnel = tunnelReq
	}

	// Record tunnel request for test assertions.
	if path == string(edgeservice.ProviderPoolPathTunnel) || path == "" {
		s.submitMu.Lock()
		s.tunnelReqs = append(s.tunnelReqs, tunnelReq)
		s.submitMu.Unlock()
	}

	if err != nil {
		return nil, err
	}

	disp := edgeservice.RunDispatch{
		RunID:         "run-pool-tunnel",
		NodeID:        "node-pool",
		ModelGroupKey: req.Run.ModelGroupKey,
		Adapter:       "selected",
		Target:        "served",
		ProviderID:    "prov-fake",
		ProviderType:  "vllm",
		ExecutionPath: string(edgeservice.ProviderPoolPathTunnel),
		QueueReason:   "dispatched",
	}
	if path == string(edgeservice.ProviderPoolPathNormalized) {
		disp.RunID = "run-pool-normalized"
	}
	if selected := s.poolSelectedCandidate; selected.ProtocolProfile != nil {
		disp.ProfileID = selected.ProfileID
		disp.ProfileDriver = selected.ProfileDriver
		disp.ProfileOperation = req.Tunnel.Operation
		disp.ProfileCapabilities = append([]string(nil), selected.ProfileCapabilities...)
		if selected.ProviderID != "" {
			disp.ProviderID = selected.ProviderID
		}
		if selected.ActualModel != "" {
			disp.Target = selected.ActualModel
		}
	}

	if path == string(edgeservice.ProviderPoolPathTunnel) || path == "" {
		frames := s.tunnelFrames
		if frames == nil && s.tunnelProviderURL != "" {
			body := req.Tunnel.Body
			if req.Tunnel.BuildBody != nil {
				target := req.Tunnel.Target
				if s.tunnelServedTarget != "" {
					target = s.tunnelServedTarget
				}
				built, err := req.Tunnel.BuildBody(target)
				if err != nil {
					return nil, err
				}
				body = built
			}
			s.submitMu.Lock()
			s.tunnelBodies = append(s.tunnelBodies, body)
			s.submitMu.Unlock()
			relayed, err := relayProviderFixtureFrames(s.tunnelProviderURL, req.Tunnel.Method, req.Tunnel.Path, body)
			if err != nil {
				return nil, err
			}
			frames = relayed
		} else if frames == nil {
			frames = staticProviderTunnelFrames(`{"ok":true}`)
			// Record the built tunnel body even when using a static frame source.
			body := req.Tunnel.Body
			if req.Tunnel.BuildBody != nil {
				target := req.Tunnel.Target
				if s.tunnelServedTarget != "" {
					target = s.tunnelServedTarget
				}
				if built, buildErr := req.Tunnel.BuildBody(target); buildErr == nil {
					body = built
				}
			}
			s.submitMu.Lock()
			s.tunnelBodies = append(s.tunnelBodies, body)
			s.submitMu.Unlock()
		} else {
			body := req.Tunnel.Body
			if req.Tunnel.BuildBody != nil {
				target := req.Tunnel.Target
				if s.tunnelServedTarget != "" {
					target = s.tunnelServedTarget
				}
				if built, err := req.Tunnel.BuildBody(target); err == nil {
					body = built
				}
			}
			s.submitMu.Lock()
			s.tunnelBodies = append(s.tunnelBodies, body)
			s.submitMu.Unlock()
		}

		handle := &fakeTunnelHandle{
			dispatch: edgeservice.RunDispatch{
				RunID:            "run-tunnel",
				NodeID:           "node-1",
				ProviderID:       disp.ProviderID,
				UsageAttribution: disp.UsageAttribution,
				ModelGroupKey:    req.Tunnel.ModelGroupKey,
				Adapter:          req.Tunnel.Adapter,
				Target:           disp.Target,
				SessionID:        req.Tunnel.SessionID,
				TimeoutSec:       5,
			},
			headers: req.Tunnel.Headers,
			frames:  frames,
		}
		// Store for later inspection by tests.
		s.submitMu.Lock()
		s.lastTunnelHandle = handle
		s.submitMu.Unlock()
		return &edgeservice.ProviderPoolDispatchResult{
			Path:         edgeservice.ProviderPoolPathTunnel,
			Tunnel:       handle,
			DispatchInfo: disp,
		}, nil
	}

	// Normalized path (mimics Ollama/CLI behavior).
	// Invoke the optional PrepareRun hook so tests can assert it only runs
	// on the normalized path (mirrors real SubmitProviderPool behavior).
	if req.PrepareRun != nil {
		s.submitMu.Lock()
		s.poolPrepareRunCalled = true
		s.submitMu.Unlock()
		prepared, prepErr := req.PrepareRun(req.Run)
		if prepErr != nil {
			return nil, prepErr
		}
		req.Run = prepared
	}
	s.submitMu.Lock()
	s.reqs = append(s.reqs, req.Run)
	s.submitMu.Unlock()

	runChan := s.poolRunFrames
	if runChan == nil {
		runChan = make(chan *iop.RunEvent, 1)
		runChan <- &iop.RunEvent{RunId: disp.RunID, Type: "complete"}
		close(runChan)
	}
	return &edgeservice.ProviderPoolDispatchResult{
		Path: edgeservice.ProviderPoolPathNormalized,
		Run: &fakeRunResult{
			dispatch: disp,
			events:   runChan,
		},
		DispatchInfo: disp,
	}, nil
}

// buildPoolResult assembles a ProviderPoolDispatchResult from a pre-configured
// custom result. It handles tunnel request recording, frame relay, and
// normalized path run construction.
func (s *providerFakeRunService) buildPoolResult(req edgeservice.ProviderPoolDispatchRequest, result edgeservice.ProviderPoolDispatchResult) (*edgeservice.ProviderPoolDispatchResult, error) {
	if result.Path == edgeservice.ProviderPoolPathTunnel || result.Path == "" {
		s.submitMu.Lock()
		s.tunnelReqs = append(s.tunnelReqs, req.Tunnel)
		s.submitMu.Unlock()
		var frames chan *iop.ProviderTunnelFrame
		if frames == nil && s.tunnelProviderURL != "" {
			body := req.Tunnel.Body
			if req.Tunnel.BuildBody != nil {
				target := req.Tunnel.Target
				if s.tunnelServedTarget != "" {
					target = s.tunnelServedTarget
				}
				if built, err := req.Tunnel.BuildBody(target); err == nil {
					body = built
				}
			}
			s.submitMu.Lock()
			s.tunnelBodies = append(s.tunnelBodies, body)
			s.submitMu.Unlock()
			relayed, err := relayProviderFixtureFrames(s.tunnelProviderURL, req.Tunnel.Method, req.Tunnel.Path, body)
			if err != nil {
				return nil, err
			}
			frames = relayed
		} else if frames == nil {
			frames = staticProviderTunnelFrames(`{"ok":true}`)
		}
		disp := result.DispatchInfo
		if disp.RunID == "" {
			disp.RunID = "run-tunnel"
		}
		if disp.NodeID == "" {
			disp.NodeID = "node-1"
		}
		handle := &fakeTunnelHandle{
			dispatch:    disp,
			headers:     req.Tunnel.Headers,
			frames:      frames,
			waitTimeout: s.tunnelWaitTimeout,
		}
		s.submitMu.Lock()
		s.lastTunnelHandle = handle
		s.submitMu.Unlock()
		return &edgeservice.ProviderPoolDispatchResult{
			Path:         edgeservice.ProviderPoolPathTunnel,
			Tunnel:       handle,
			DispatchInfo: disp,
		}, nil
	}

	// Normalized path.
	disp := result.DispatchInfo
	if disp.RunID == "" {
		disp.RunID = "run-pool-normalized"
	}
	if disp.NodeID == "" {
		disp.NodeID = "node-pool"
	}
	// Invoke the optional PrepareRun hook on the normalized path (mirrors
	// real SubmitProviderPool behavior).
	if req.PrepareRun != nil {
		prepared, prepErr := req.PrepareRun(req.Run)
		if prepErr != nil {
			return nil, prepErr
		}
		req.Run = prepared
	}
	s.submitMu.Lock()
	s.poolPrepareRunCalled = true
	s.submitMu.Unlock()
	s.submitMu.Lock()
	s.reqs = append(s.reqs, req.Run)
	s.submitMu.Unlock()
	// Use the Run from the custom result if provided, otherwise create a
	// default completing run.
	if result.Run != nil {
		return &edgeservice.ProviderPoolDispatchResult{
			Path:         edgeservice.ProviderPoolPathNormalized,
			Run:          result.Run,
			DispatchInfo: disp,
		}, nil
	}
	runChan := make(chan *iop.RunEvent, 1)
	runChan <- &iop.RunEvent{RunId: disp.RunID, Type: "complete"}
	close(runChan)
	return &edgeservice.ProviderPoolDispatchResult{
		Path: edgeservice.ProviderPoolPathNormalized,
		Run: &fakeRunResult{
			dispatch: disp,
			events:   runChan,
		},
		DispatchInfo: disp,
	}, nil
}

// tunnelHandleSnapshot returns the tunnel handle from the last SubmitProviderPool
// call, allowing tests to check headers set by the handler.
type tunnelHandleSnapshot struct {
	headers map[string]string
	runID   string
	nodeID  string
}

func (s *providerFakeRunService) tunnelHandleSnapshot() tunnelHandleSnapshot {
	s.submitMu.Lock()
	defer s.submitMu.Unlock()
	h := s.lastTunnelHandle
	if h == nil {
		return tunnelHandleSnapshot{}
	}
	return tunnelHandleSnapshot{
		headers: h.Headers(),
		runID:   h.Dispatch().RunID,
		nodeID:  h.Dispatch().NodeID,
	}
}

// poolSubmitCountSnapshot returns the number of times SubmitProviderPool was
// called. Tests assert exactly-once dispatch from the handler.
func (s *providerFakeRunService) poolSubmitCountSnapshot() int {
	s.submitMu.Lock()
	defer s.submitMu.Unlock()
	return s.poolSubmitCount
}

// poolLastRunSnapshot returns a copy of the last SubmitRunRequest passed to
// SubmitProviderPool. Tests assert ModelGroupKey is non-empty to catch the
// "empty Run" regression.
func (s *providerFakeRunService) poolLastRunSnapshot() edgeservice.SubmitRunRequest {
	s.submitMu.Lock()
	defer s.submitMu.Unlock()
	return s.poolLastRun
}

// poolPrepareRunCalledSnapshot returns whether the PrepareRun hook was
// invoked by the last SubmitProviderPool call. Only meaningful when the last
// call was on the normalized path.
func (s *providerFakeRunService) poolPrepareRunCalledSnapshot() bool {
	s.submitMu.Lock()
	defer s.submitMu.Unlock()
	return s.poolPrepareRunCalled
}

func (s *providerFakeRunService) tunnelReqsSnapshot() []edgeservice.SubmitProviderTunnelRequest {
	s.submitMu.Lock()
	reqs := make([]edgeservice.SubmitProviderTunnelRequest, len(s.tunnelReqs))
	copy(reqs, s.tunnelReqs)
	s.submitMu.Unlock()
	return reqs
}

func (s *providerFakeRunService) tunnelBodiesSnapshot() [][]byte {
	s.submitMu.Lock()
	defer s.submitMu.Unlock()
	return append([][]byte(nil), s.tunnelBodies...)
}

// --- scripted provider-pool admission double (S16 host fixtures) -------------

// scriptedPoolAttempt describes one SubmitProviderPool admission the scripted
// double must produce, in order. It carries the actual selection result the
// real pool would compute (path, provider, model), so a fixture can assert that
// one request runtime observed a real actual-path switch.
type scriptedPoolAttempt struct {
	// path is the execution path this admission resolves to, as the
	// service-layer string value ("normalized" or "provider_tunnel").
	path      string
	runID     string
	provider  string
	target    string
	runEvents chan *iop.RunEvent
	frames    chan *iop.ProviderTunnelFrame
	err       error
	// prepareErr, when set, is returned by the recorded PrepareRun hook so a
	// fixture can exercise a normalized candidate whose validation fails.
	prepareErr error
}

// scriptedPoolRunService is a thread-safe runService double whose provider-pool
// admissions are scripted in order. It records every dispatch, transport close,
// and cancel so an S16 fixture can assert single-dispatch, exactly-once
// cleanup, and no leak from a discarded attempt.
type scriptedPoolRunService struct {
	mu       sync.Mutex
	attempts []scriptedPoolAttempt
	next     int

	poolRequests   []edgeservice.ProviderPoolDispatchRequest
	runRequests    []edgeservice.SubmitRunRequest
	tunnelRequests []edgeservice.SubmitProviderTunnelRequest
	cancels        []edgeservice.CancelRunRequest
	runCloses      []string
	tunnelCloses   []string
	prepareTunnel  int
	prepareRun     int
}

func newScriptedPoolRunService(attempts ...scriptedPoolAttempt) *scriptedPoolRunService {
	return &scriptedPoolRunService{attempts: attempts}
}

func (s *scriptedPoolRunService) SubmitRun(_ context.Context, req edgeservice.SubmitRunRequest) (edgeservice.RunResult, error) {
	s.mu.Lock()
	s.runRequests = append(s.runRequests, req)
	s.mu.Unlock()
	return nil, fmt.Errorf("scriptedPoolRunService: direct SubmitRun is not part of a provider-pool fixture")
}

func (s *scriptedPoolRunService) SubmitProviderTunnel(_ context.Context, req edgeservice.SubmitProviderTunnelRequest) (edgeservice.ProviderTunnelResult, error) {
	s.mu.Lock()
	s.tunnelRequests = append(s.tunnelRequests, req)
	s.mu.Unlock()
	return nil, fmt.Errorf("scriptedPoolRunService: direct SubmitProviderTunnel is not part of a provider-pool fixture")
}

func (s *scriptedPoolRunService) OllamaAPI(context.Context, edgeservice.OllamaAPIRequest) (edgeservice.OllamaAPIView, error) {
	return edgeservice.OllamaAPIView{StatusCode: http.StatusOK}, nil
}

func (s *scriptedPoolRunService) CancelRun(_ context.Context, req edgeservice.CancelRunRequest) (edgeservice.CommandResult, error) {
	s.mu.Lock()
	s.cancels = append(s.cancels, req)
	s.mu.Unlock()
	return edgeservice.CommandResult{NodeID: req.NodeRef}, nil
}

func (s *scriptedPoolRunService) SubmitProviderPool(_ context.Context, req edgeservice.ProviderPoolDispatchRequest) (*edgeservice.ProviderPoolDispatchResult, error) {
	s.mu.Lock()
	s.poolRequests = append(s.poolRequests, req)
	if s.next >= len(s.attempts) {
		s.mu.Unlock()
		return nil, fmt.Errorf("scriptedPoolRunService: no scripted attempt left")
	}
	attempt := s.attempts[s.next]
	s.next++
	s.mu.Unlock()

	if attempt.err != nil {
		return nil, attempt.err
	}

	dispatch := edgeservice.RunDispatch{
		RunID:         attempt.runID,
		NodeID:        "node-pool",
		ModelGroupKey: req.Run.ModelGroupKey,
		Adapter:       "selected",
		Target:        attempt.target,
		ProviderID:    attempt.provider,
		ProviderType:  "scripted",
		ExecutionPath: string(attempt.path),
		SessionID:     req.Run.SessionID,
		TimeoutSec:    5,
	}

	if attempt.path == string(edgeservice.ProviderPoolPathTunnel) {
		tunnelReq := req.Tunnel
		if req.PrepareProtocolTunnel != nil {
			s.mu.Lock()
			s.prepareTunnel++
			s.mu.Unlock()
			prepared, err := req.PrepareProtocolTunnel(tunnelReq, edgeservice.ProviderPoolCandidate{
				ActualModel: attempt.target, ProviderID: attempt.provider, ExecutionPath: string(attempt.path),
			})
			if err != nil {
				return nil, err
			}
			tunnelReq = prepared
		} else if req.PrepareTunnel != nil {
			s.mu.Lock()
			s.prepareTunnel++
			s.mu.Unlock()
			prepared, err := req.PrepareTunnel(tunnelReq)
			if err != nil {
				return nil, err
			}
			tunnelReq = prepared
		}
		if tunnelReq.BuildBody != nil {
			if _, err := tunnelReq.BuildBody(attempt.target); err != nil {
				return nil, err
			}
		}
		s.mu.Lock()
		s.tunnelRequests = append(s.tunnelRequests, tunnelReq)
		s.mu.Unlock()
		return &edgeservice.ProviderPoolDispatchResult{
			Path:         edgeservice.ProviderPoolPathTunnel,
			Tunnel:       &countingTunnelHandle{service: s, dispatch: dispatch, frames: attempt.frames},
			DispatchInfo: dispatch,
		}, nil
	}

	runReq := req.Run
	if req.PrepareRun != nil {
		s.mu.Lock()
		s.prepareRun++
		s.mu.Unlock()
		if attempt.prepareErr != nil {
			return nil, attempt.prepareErr
		}
		prepared, err := req.PrepareRun(runReq)
		if err != nil {
			return nil, err
		}
		runReq = prepared
	}
	s.mu.Lock()
	s.runRequests = append(s.runRequests, runReq)
	s.mu.Unlock()
	return &edgeservice.ProviderPoolDispatchResult{
		Path:         edgeservice.ProviderPoolPathNormalized,
		Run:          &countingRunHandle{service: s, dispatch: dispatch, events: attempt.runEvents},
		DispatchInfo: dispatch,
	}, nil
}

func (s *scriptedPoolRunService) snapshot() (pools int, cancels []string, runCloses, tunnelCloses []string, runReqs []edgeservice.SubmitRunRequest, tunnelReqs []edgeservice.SubmitProviderTunnelRequest) {
	s.mu.Lock()
	defer s.mu.Unlock()
	ids := make([]string, 0, len(s.cancels))
	for _, c := range s.cancels {
		ids = append(ids, c.RunID)
	}
	return len(s.poolRequests),
		ids,
		append([]string(nil), s.runCloses...),
		append([]string(nil), s.tunnelCloses...),
		append([]edgeservice.SubmitRunRequest(nil), s.runRequests...),
		append([]edgeservice.SubmitProviderTunnelRequest(nil), s.tunnelRequests...)
}

func (s *scriptedPoolRunService) poolSubmits() int {
	s.mu.Lock()
	defer s.mu.Unlock()
	return len(s.poolRequests)
}

type countingRunHandle struct {
	service  *scriptedPoolRunService
	dispatch edgeservice.RunDispatch
	events   chan *iop.RunEvent
}

func (h *countingRunHandle) Dispatch() edgeservice.RunDispatch { return h.dispatch }
func (h *countingRunHandle) Stream() edgeservice.RunStream {
	return edgeservice.RunStream{Events: h.events}
}
func (h *countingRunHandle) Close() {
	h.service.mu.Lock()
	h.service.runCloses = append(h.service.runCloses, h.dispatch.RunID)
	h.service.mu.Unlock()
}
func (h *countingRunHandle) WaitTimeout() time.Duration { return 5 * time.Second }

type countingTunnelHandle struct {
	service  *scriptedPoolRunService
	dispatch edgeservice.RunDispatch
	headers  map[string]string
	frames   chan *iop.ProviderTunnelFrame
}

func (h *countingTunnelHandle) Headers() map[string]string        { return h.headers }
func (h *countingTunnelHandle) SetHeaders(hdrs map[string]string) { h.headers = hdrs }
func (h *countingTunnelHandle) Dispatch() edgeservice.RunDispatch { return h.dispatch }
func (h *countingTunnelHandle) Stream() edgeservice.ProviderTunnelStream {
	return edgeservice.ProviderTunnelStream{Frames: h.frames}
}
func (h *countingTunnelHandle) Close() {
	h.service.mu.Lock()
	h.service.tunnelCloses = append(h.service.tunnelCloses, h.dispatch.RunID)
	h.service.mu.Unlock()
}
func (h *countingTunnelHandle) WaitTimeout() time.Duration { return 5 * time.Second }

var (
	_ runService                       = (*scriptedPoolRunService)(nil)
	_ edgeservice.RunResult            = (*countingRunHandle)(nil)
	_ edgeservice.ProviderTunnelResult = (*countingTunnelHandle)(nil)
)

// bufferedTunnelFrames builds a closed, fully buffered tunnel frame channel.
func bufferedTunnelFrames(frames ...*iop.ProviderTunnelFrame) chan *iop.ProviderTunnelFrame {
	ch := make(chan *iop.ProviderTunnelFrame, len(frames))
	for _, f := range frames {
		ch <- f
	}
	close(ch)
	return ch
}

// relayProviderFixtureFrames performs the provider HTTP request the way the
// Node relay does and converts status/header/body bytes into ordered tunnel
// frames, chunking the body at odd boundaries to exercise ordered writes.
func relayProviderFixtureFrames(providerURL, method, path string, body []byte) (chan *iop.ProviderTunnelFrame, error) {
	httpReq, err := http.NewRequest(method, providerURL+path, bytes.NewReader(body))
	if err != nil {
		return nil, err
	}
	httpReq.Header.Set("Content-Type", "application/json")
	resp, err := http.DefaultClient.Do(httpReq)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	respBody, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}
	headers := make(map[string]string)
	for k, values := range resp.Header {
		if len(values) > 0 {
			headers[k] = values[0]
		}
	}

	const chunkSize = 7
	frames := make(chan *iop.ProviderTunnelFrame, len(respBody)/chunkSize+3)
	seq := int64(0)
	frames <- &iop.ProviderTunnelFrame{
		RunId:      "run-tunnel",
		TunnelId:   "tunnel-1",
		Sequence:   seq,
		Kind:       iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_RESPONSE_START,
		StatusCode: int32(resp.StatusCode),
		Headers:    headers,
	}
	seq++
	for off := 0; off < len(respBody); off += chunkSize {
		end := off + chunkSize
		if end > len(respBody) {
			end = len(respBody)
		}
		frames <- &iop.ProviderTunnelFrame{
			RunId:    "run-tunnel",
			TunnelId: "tunnel-1",
			Sequence: seq,
			Kind:     iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_BODY,
			Body:     append([]byte(nil), respBody[off:end]...),
		}
		seq++
	}
	frames <- &iop.ProviderTunnelFrame{
		RunId:    "run-tunnel",
		TunnelId: "tunnel-1",
		Sequence: seq,
		Kind:     iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_END,
		End:      true,
	}
	close(frames)
	return frames, nil
}

func staticProviderTunnelFrames(body string) chan *iop.ProviderTunnelFrame {
	frames := make(chan *iop.ProviderTunnelFrame, 3)
	frames <- &iop.ProviderTunnelFrame{Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_RESPONSE_START, StatusCode: 200}
	frames <- &iop.ProviderTunnelFrame{Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_BODY, Body: []byte(body)}
	frames <- &iop.ProviderTunnelFrame{Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_END, End: true}
	close(frames)
	return frames
}