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 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.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 }