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

package openai

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"net"
	"strings"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	toki "git.toki-labs.com/toki/proto-socket/go"
	edgenode "iop/apps/edge/internal/node"
	edgeservice "iop/apps/edge/internal/service"
	"iop/packages/go/config"
	iop "iop/proto/gen/iop"
)

func newTestServiceHarness(t *testing.T, executor *SingleRequestExecutor) (*edgeservice.Service, *edgeservice.SingleRequestBinding, *workNodeHarness) {
	t.Helper()

	nodeHarness := newWorkNodeHarness()
	registry := edgenode.NewRegistry()
	store := edgenode.NewNodeStore()

	for _, pair := range []struct{ nodeID, workspaceRef string }{
		{"node", "workspace"},
		{"node-0", "workspace-0"},
		{"node-1", "workspace-1"},
	} {
		edgeConn, nodeConn := net.Pipe()
		edgeClient := toki.NewTcpClient(edgeConn, 0, 0, toki.ParserMap{
			toki.TypeNameOf(&iop.WorkspaceOpenResponse{}):     parseAnthropicWorkspaceMessage(&iop.WorkspaceOpenResponse{}),
			toki.TypeNameOf(&iop.WorkspaceArtifactResponse{}): parseAnthropicWorkspaceMessage(&iop.WorkspaceArtifactResponse{}),
			toki.TypeNameOf(&iop.WorkspaceToolResponse{}):     parseAnthropicWorkspaceMessage(&iop.WorkspaceToolResponse{}),
			toki.TypeNameOf(&iop.WorkspaceCancelResponse{}):   parseAnthropicWorkspaceMessage(&iop.WorkspaceCancelResponse{}),
			toki.TypeNameOf(&iop.WorkspaceCleanupResponse{}):  parseAnthropicWorkspaceMessage(&iop.WorkspaceCleanupResponse{}),
		})
		nodeClient := toki.NewTcpClient(nodeConn, 0, 0, toki.ParserMap{
			toki.TypeNameOf(&iop.WorkspaceOpenRequest{}):     parseAnthropicWorkspaceMessage(&iop.WorkspaceOpenRequest{}),
			toki.TypeNameOf(&iop.WorkspaceArtifactRequest{}): parseAnthropicWorkspaceMessage(&iop.WorkspaceArtifactRequest{}),
			toki.TypeNameOf(&iop.WorkspaceToolRequest{}):     parseAnthropicWorkspaceMessage(&iop.WorkspaceToolRequest{}),
			toki.TypeNameOf(&iop.WorkspaceCancelRequest{}):   parseAnthropicWorkspaceMessage(&iop.WorkspaceCancelRequest{}),
			toki.TypeNameOf(&iop.WorkspaceCleanupRequest{}):  parseAnthropicWorkspaceMessage(&iop.WorkspaceCleanupRequest{}),
		})
		t.Cleanup(func() {
			_ = edgeClient.Close()
			_ = nodeClient.Close()
		})

		registry.Register(&edgenode.NodeEntry{NodeID: pair.nodeID, Alias: "work-node", Client: edgeClient})
		store.Add(&edgenode.NodeRecord{ID: pair.nodeID, Alias: "work-node", Token: "work-node-token", Workspaces: []config.WorkspaceDefinition{{
			Ref: pair.workspaceRef, Platform: "darwin", Root: "/Users/operator/project",
			Operations: []config.WorkspaceOperation{config.WorkspaceOpRead, config.WorkspaceOpWrite, config.WorkspaceOpCommand},
			Commands:   []config.WorkspaceCommandDefinition{{ID: "verify", Executable: "/usr/bin/true"}}, EnvironmentAllowlist: []string{"SAFE"},
			MaxReadBytes: 4096, MaxWriteBytes: 4096, MaxOutputBytes: 4096, MaxCommandTimeoutMS: 30000,
		}}})

		nodeHarness.install(nodeClient)
	}

	d := validDispatch()
	planBinding := edgeservice.SingleRequestStageBinding{Model: d.ModelGroupKey, Options: map[string]any{"reasoning_effort": "high"}, Dispatch: &d}
	workBinding := edgeservice.SingleRequestStageBinding{Model: d.ModelGroupKey, Options: map[string]any{"temperature": 0.2}, Dispatch: &d}
	reviewBinding := edgeservice.SingleRequestStageBinding{Model: d.ModelGroupKey, Options: map[string]any{"reasoning_effort": "high"}, Dispatch: &d}

	limits := validLimits()
	limits.WallClockMS = 60000
	limits.StageTimeoutMS = 30000
	limits.MaxToolIterations = 4
	binding, err := edgeservice.NewSingleRequestBinding("public", "workspace", planBinding, workBinding, reviewBinding, limits)
	if err != nil {
		t.Fatal(err)
	}

	service := edgeservice.New(registry, nil)
	service.SetNodeStore(store)
	service.SetSingleRequestExecutor(executor)

	return service, binding, nodeHarness
}

func waitExecutionResult(exec edgeservice.SingleRequestExecution) (edgeservice.SingleRequestResult, error) {
	for prog := range exec.Progress() {
		if prog.Stage == edgeservice.SingleRequestStateFinalizing {
			_ = exec.AcknowledgeTerminal(true)
		}
	}
	return exec.Wait()
}

func executorPlanBody(plan, verification string) []byte {
	b, _ := json.Marshal(map[string]any{
		"plan":         plan,
		"verification": verification,
	})
	return successBodyWithThoughtSignature(string(b))
}

func executorWorkBody(completion, verification string) []byte {
	b, _ := json.Marshal(map[string]any{
		"completion":   completion,
		"verification": verification,
	})
	return successBody(string(b))
}

func executorReviewPassBody(output, summary string) []byte {
	b, _ := json.Marshal(map[string]any{
		"decision": "pass",
		"output":   output,
		"summary":  summary,
	})
	return successBodyWithThoughtSignature(string(b))
}

func TestSingleRequestExecutorInterface(t *testing.T) {
	executor := NewSingleRequestExecutor(&mockService{})
	var _ edgeservice.SingleRequestExecutor = executor
	var _ edgeservice.SingleRequestToolContinuation = executor
}

func TestSingleRequestExecutorPass(t *testing.T) {
	responses := [][]byte{
		executorPlanBody("Execute step 1", "Verify step 1"),
		workToolBody("work-pass-1", edgeservice.InternalWorkspaceToolRead, `{"relative_path":"output.txt"}`),
		executorWorkBody("Completed work step 1", "Verified work step 1"),
		executorReviewPassBody("Final Approved Output", "Review passed cleanly"),
	}

	var callIndex atomic.Int32
	mockSvc := &mockService{
		submit: func(_ context.Context, req edgeservice.ProviderPoolDispatchRequest) (*edgeservice.ProviderPoolDispatchResult, error) {
			idx := int(callIndex.Add(1) - 1)
			if idx >= len(responses) {
				t.Fatalf("unexpected call index %d", idx)
			}
			return &edgeservice.ProviderPoolDispatchResult{
				Path:         edgeservice.ProviderPoolPathTunnel,
				Tunnel:       &mockTunnel{frames: framesFor(responses[idx])},
				DispatchInfo: matchingDispatch(),
			}, nil
		},
	}

	executor := NewSingleRequestExecutor(mockSvc)
	svc, binding, _ := newTestServiceHarness(t, executor)

	req := edgeservice.SingleRequestRequest{
		RequestID: "req-pass-1",
		Binding:   binding,
		Prompt:    "Complete assignment",
	}

	exec, err := svc.StartSingleRequest(context.Background(), req)
	if err != nil {
		t.Fatal(err)
	}

	result, err := waitExecutionResult(exec)
	if err != nil {
		t.Fatalf("unexpected wait error: %v", err)
	}

	if result.Output != "Final Approved Output" {
		t.Fatalf("got output %q, want %q", result.Output, "Final Approved Output")
	}

	if exec.State() != edgeservice.SingleRequestStateCompleted {
		t.Fatalf("state = %s, want completed", exec.State())
	}
}

func TestSingleRequestExecutorInspection(t *testing.T) {
	responses := [][]byte{
		executorPlanBody("Plan inspect", "Verify plan"),
		workToolBody("work-inspect-1", edgeservice.InternalWorkspaceToolRead, `{"relative_path":"output.txt"}`),
		executorWorkBody("Work done", "Work verified"),
		workToolBody("tool-inspect-1", edgeservice.InternalWorkspaceToolRead, `{"relative_path":"output.txt"}`),
		executorReviewPassBody("Inspected Final Output", "Inspection approved"),
	}

	var callIndex atomic.Int32
	mockSvc := &mockService{
		submit: func(_ context.Context, req edgeservice.ProviderPoolDispatchRequest) (*edgeservice.ProviderPoolDispatchResult, error) {
			idx := int(callIndex.Add(1) - 1)
			if idx >= len(responses) {
				t.Fatalf("unexpected call index %d", idx)
			}
			return &edgeservice.ProviderPoolDispatchResult{
				Path:         edgeservice.ProviderPoolPathTunnel,
				Tunnel:       &mockTunnel{frames: framesFor(responses[idx])},
				DispatchInfo: matchingDispatch(),
			}, nil
		},
	}

	executor := NewSingleRequestExecutor(mockSvc)
	svc, binding, _ := newTestServiceHarness(t, executor)

	req := edgeservice.SingleRequestRequest{
		RequestID: "req-inspect-1",
		Binding:   binding,
		Prompt:    "Inspect output",
	}

	exec, err := svc.StartSingleRequest(context.Background(), req)
	if err != nil {
		t.Fatal(err)
	}

	result, err := waitExecutionResult(exec)
	if err != nil {
		t.Fatalf("unexpected wait error: %v", err)
	}

	if result.Output != "Inspected Final Output" {
		t.Fatalf("got output %q, want %q", result.Output, "Inspected Final Output")
	}
}

func TestSingleRequestExecutorRepair(t *testing.T) {
	responses := [][]byte{
		executorPlanBody("Plan repair", "Verify plan"),
		workToolBody("work-repair-1", edgeservice.InternalWorkspaceToolRead, `{"relative_path":"output.txt"}`),
		executorWorkBody("Work initial", "Work initial verify"),
		workToolBody("tool-repair-1", edgeservice.InternalWorkspaceToolWrite, `{"relative_path":"output.txt","content":"fixed content"}`),
		executorReviewPassBody("Repaired Final Output", "Repair approved"),
	}

	var callIndex atomic.Int32
	mockSvc := &mockService{
		submit: func(_ context.Context, req edgeservice.ProviderPoolDispatchRequest) (*edgeservice.ProviderPoolDispatchResult, error) {
			idx := int(callIndex.Add(1) - 1)
			if idx >= len(responses) {
				t.Fatalf("unexpected call index %d", idx)
			}
			return &edgeservice.ProviderPoolDispatchResult{
				Path:         edgeservice.ProviderPoolPathTunnel,
				Tunnel:       &mockTunnel{frames: framesFor(responses[idx])},
				DispatchInfo: matchingDispatch(),
			}, nil
		},
	}

	executor := NewSingleRequestExecutor(mockSvc)
	svc, binding, _ := newTestServiceHarness(t, executor)

	req := edgeservice.SingleRequestRequest{
		RequestID: "req-repair-1",
		Binding:   binding,
		Prompt:    "Repair output",
	}

	exec, err := svc.StartSingleRequest(context.Background(), req)
	if err != nil {
		t.Fatal(err)
	}

	result, err := waitExecutionResult(exec)
	if err != nil {
		t.Fatalf("unexpected wait error: %v", err)
	}

	if result.Output != "Repaired Final Output" {
		t.Fatalf("got output %q, want %q", result.Output, "Repaired Final Output")
	}
}

func TestSingleRequestExecutorConcurrentToolIsolation(t *testing.T) {
	const concurrency = 2

	toolArrived := make(chan string, concurrency)
	releaseToolResponses := make(chan struct{})

	mockSvc := &mockService{
		submit: func(_ context.Context, req edgeservice.ProviderPoolDispatchRequest) (*edgeservice.ProviderPoolDispatchResult, error) {
			reqID := req.Tunnel.SessionID
			if reqID == "" {
				return nil, errors.New("missing session ID in tunnel request")
			}

			reqBody, err := req.Tunnel.BuildBody("gemini-3.6-flash")
			if err != nil {
				return nil, err
			}
			bodyStr := string(reqBody)

			var resp []byte
			if strings.Contains(bodyStr, "Produce exactly one JSON object with non-empty string fields plan") {
				resp = executorPlanBody(fmt.Sprintf("Plan for %s", reqID), fmt.Sprintf("Verify plan for %s", reqID))
			} else if strings.Contains(bodyStr, "Read the supplied plan") {
				if !strings.Contains(bodyStr, "typed-result-") {
					resp = workToolBody("colliding-tool-id", edgeservice.InternalWorkspaceToolRead, fmt.Sprintf(`{"relative_path":"output-%s.txt"}`, reqID))
				} else {
					wantPlan := fmt.Sprintf("Plan for %s", reqID)
					wantResult := fmt.Sprintf("typed-result-%s", reqID)
					if !strings.Contains(bodyStr, wantPlan) || !strings.Contains(bodyStr, wantResult) {
						return nil, fmt.Errorf("isolation failure for %s: missing expected plan/result in body: %s", reqID, bodyStr)
					}
					for otherIdx := 0; otherIdx < concurrency; otherIdx++ {
						otherID := fmt.Sprintf("req-iso-%d", otherIdx)
						if otherID != reqID {
							otherPlan := fmt.Sprintf("Plan for %s", otherID)
							otherResult := fmt.Sprintf("typed-result-%s", otherID)
							if strings.Contains(bodyStr, otherPlan) || strings.Contains(bodyStr, otherResult) {
								return nil, fmt.Errorf("isolation failure for %s: body contains data from %s", reqID, otherID)
							}
						}
					}
					resp = executorWorkBody(fmt.Sprintf("Work completion for %s", reqID), fmt.Sprintf("Work verify for %s", reqID))
				}
			} else {
				wantWorkComp := fmt.Sprintf("Work completion for %s", reqID)
				if !strings.Contains(bodyStr, wantWorkComp) {
					return nil, fmt.Errorf("isolation failure for %s: review body missing work completion: %s", reqID, bodyStr)
				}
				resp = executorReviewPassBody(fmt.Sprintf("Reviewer Approved for %s", reqID), fmt.Sprintf("Review pass for %s", reqID))
			}

			return &edgeservice.ProviderPoolDispatchResult{
				Path:         edgeservice.ProviderPoolPathTunnel,
				Tunnel:       &mockTunnel{frames: framesFor(resp)},
				DispatchInfo: matchingDispatch(),
			}, nil
		},
	}

	executor := NewSingleRequestExecutor(mockSvc)
	svc, binding, nodeHarness := newTestServiceHarness(t, executor)

	nodeHarness.toolResponder = func(req *iop.WorkspaceToolRequest) *iop.WorkspaceToolResponse {
		toolArrived <- req.GetRequestId()
		<-releaseToolResponses
		return &iop.WorkspaceToolResponse{
			RequestId:  req.GetRequestId(),
			StageId:    req.GetStageId(),
			ToolCallId: req.GetToolCallId(),
			Status:     iop.WorkspaceStatus_WORKSPACE_STATUS_SUCCESS,
			Content:    []byte("typed-result-" + req.GetRequestId()),
		}
	}

	expectedByRequest := make(map[string]string, concurrency)
	for i := 0; i < concurrency; i++ {
		reqID := fmt.Sprintf("req-iso-%d", i)
		expectedByRequest[reqID] = fmt.Sprintf("Reviewer Approved for %s", reqID)
	}

	var wg sync.WaitGroup
	wg.Add(concurrency)

	for i := 0; i < concurrency; i++ {
		go func(id int) {
			defer wg.Done()

			reqID := fmt.Sprintf("req-iso-%d", id)
			reqBinding, err := edgeservice.NewSingleRequestBinding("public", fmt.Sprintf("workspace-%d", id), binding.Plan, binding.Work, binding.Review, binding.Limits)
			if err != nil {
				t.Errorf("request %s binding error: %v", reqID, err)
				return
			}
			req := edgeservice.SingleRequestRequest{
				RequestID: reqID,
				Binding:   reqBinding,
				Prompt:    fmt.Sprintf("Task prompt for %s", reqID),
			}

			exec, err := svc.StartSingleRequest(context.Background(), req)
			if err != nil {
				t.Errorf("request %s start error: %v", reqID, err)
				return
			}

			result, err := waitExecutionResult(exec)
			if err != nil {
				t.Errorf("request %s wait error: %v", reqID, err)
				return
			}

			if result.Output != expectedByRequest[reqID] {
				t.Errorf("request %s output = %q, want %q", reqID, result.Output, expectedByRequest[reqID])
			}
		}(i)
	}

	for i := 0; i < concurrency; i++ {
		select {
		case <-toolArrived:
		case <-time.After(5 * time.Second):
			t.Fatal("timed out waiting for tool arrival")
		}
	}

	if got := executor.bridge.pendingCount(); got != concurrency {
		t.Fatalf("bridge pending count = %d, want %d before response release", got, concurrency)
	}

	close(releaseToolResponses)

	wg.Wait()

	if len(nodeHarness.toolRequestsByRequest) != concurrency {
		t.Fatalf("toolRequestsByRequest count = %d, want %d", len(nodeHarness.toolRequestsByRequest), concurrency)
	}
	if len(nodeHarness.toolResponsesByRequest) != concurrency {
		t.Fatalf("toolResponsesByRequest count = %d, want %d", len(nodeHarness.toolResponsesByRequest), concurrency)
	}

	for i := 0; i < concurrency; i++ {
		reqID := fmt.Sprintf("req-iso-%d", i)

		gotPlan := string(nodeHarness.plansByRequest[reqID])
		wantPlan := fmt.Sprintf("Plan for %s", reqID)
		if !strings.Contains(gotPlan, wantPlan) {
			t.Errorf("request %s plan = %q, want containing %q", reqID, gotPlan, wantPlan)
		}

		reqs := nodeHarness.toolRequestsByRequest[reqID]
		if len(reqs) != 1 {
			t.Errorf("request %s tool requests count = %d, want 1", reqID, len(reqs))
		} else if reqs[0].GetToolCallId() != "colliding-tool-id" {
			t.Errorf("request %s tool call ID = %q, want colliding-tool-id", reqID, reqs[0].GetToolCallId())
		}

		resps := nodeHarness.toolResponsesByRequest[reqID]
		if len(resps) != 1 {
			t.Errorf("request %s tool responses count = %d, want 1", reqID, len(resps))
		} else {
			gotResult := string(resps[0].GetContent())
			wantResult := fmt.Sprintf("typed-result-%s", reqID)
			if gotResult != wantResult {
				t.Errorf("request %s result = %q, want %q", reqID, gotResult, wantResult)
			}
		}
	}

	if executor.bridge.pendingCount() != 0 {
		t.Fatalf("bridge pending count = %d, want 0 after concurrent completion", executor.bridge.pendingCount())
	}
}

func TestSingleRequestExecutorCancellation(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	mockSvc := &mockService{
		submit: func(_ context.Context, _ edgeservice.ProviderPoolDispatchRequest) (*edgeservice.ProviderPoolDispatchResult, error) {
			cancel()
			return nil, errors.New("cancelled in provider")
		},
	}

	executor := NewSingleRequestExecutor(mockSvc)
	svc, binding, _ := newTestServiceHarness(t, executor)
	req := edgeservice.SingleRequestRequest{
		RequestID: "req-cancel-1",
		Binding:   binding,
		Prompt:    "Cancel task",
	}

	exec, err := svc.StartSingleRequest(ctx, req)
	if err != nil {
		t.Fatal(err)
	}

	_, err = waitExecutionResult(exec)
	if !errors.Is(err, edgeservice.ErrSingleRequestCancelled) && !errors.Is(err, context.Canceled) {
		t.Fatalf("expected cancel error, got: %v", err)
	}

	if executor.bridge.pendingCount() != 0 {
		t.Fatalf("pending count = %d, want 0", executor.bridge.pendingCount())
	}
}

func TestSingleRequestExecutorParentContextOwnership(t *testing.T) {
	t.Run("cancelled parent", func(t *testing.T) {
		ctx, cancel := context.WithCancel(context.Background())
		cancel()
		controller := &qualityGateController{}
		err := submitSingleRequestClosedTerminal(
			ctx,
			"quality-request",
			controller,
			newSingleRequestQualityGate().providerFailure(ctx, context.Canceled, errProviderStageGeneric),
		)
		if !errors.Is(err, context.Canceled) || len(controller.envelopes) != 0 {
			t.Fatalf("submit error=%v envelopes=%+v, want parent cancellation and no competing terminal", err, controller.envelopes)
		}
	})

	t.Run("expired parent", func(t *testing.T) {
		ctx, cancel := context.WithDeadline(context.Background(), time.Now().Add(-time.Second))
		defer cancel()
		controller := &qualityGateController{}
		err := submitSingleRequestClosedTerminal(
			ctx,
			"quality-request",
			controller,
			newSingleRequestQualityGate().providerFailure(ctx, context.DeadlineExceeded, errProviderStageGeneric),
		)
		if !errors.Is(err, context.DeadlineExceeded) || len(controller.envelopes) != 0 {
			t.Fatalf("submit error=%v envelopes=%+v, want parent deadline and no competing terminal", err, controller.envelopes)
		}
	})

	t.Run("live parent raw cancellation", func(t *testing.T) {
		controller := &qualityGateController{}
		if err := submitSingleRequestClosedTerminal(context.Background(), "quality-request", controller, context.Canceled); err != nil {
			t.Fatalf("submit terminal: %v", err)
		}
		want := edgeservice.SingleRequestTerminalDisposition{Kind: edgeservice.SingleRequestTerminalError, ErrorClass: edgeservice.SingleRequestTerminalErrorProvider}
		if len(controller.envelopes) != 1 || controller.envelopes[0].Terminal == nil || *controller.envelopes[0].Terminal != want {
			t.Fatalf("envelopes=%+v, want one provider terminal", controller.envelopes)
		}
	})
}

func TestSingleRequestExecutorRequestBudgetOwnership(t *testing.T) {
	t.Run("request wall clock", func(t *testing.T) {
		for iteration := 0; iteration < 20; iteration++ {
			var providerCalls atomic.Int32
			mockSvc := &mockService{submit: func(ctx context.Context, _ edgeservice.ProviderPoolDispatchRequest) (*edgeservice.ProviderPoolDispatchResult, error) {
				providerCalls.Add(1)
				<-ctx.Done()
				return nil, ctx.Err()
			}}
			executor := NewSingleRequestExecutor(mockSvc)
			svc, binding, node := newTestServiceHarness(t, executor)
			binding.Limits.WallClockMS = 10
			binding.Limits.StageTimeoutMS = 10

			execution, err := svc.StartSingleRequest(context.Background(), edgeservice.SingleRequestRequest{
				RequestID: fmt.Sprintf("req-budget-%d", iteration),
				Binding:   binding,
				Prompt:    "request budget ownership",
			})
			if err != nil {
				t.Fatal(err)
			}
			terminalCount := 0
			var terminal edgeservice.SingleRequestTerminalDisposition
			for progress := range execution.Progress() {
				if progress.Terminal != nil {
					terminalCount++
					terminal = *progress.Terminal
				}
			}
			_, waitErr := execution.Wait()
			want := edgeservice.SingleRequestTerminalDisposition{Kind: edgeservice.SingleRequestTerminalError, ErrorClass: edgeservice.SingleRequestTerminalErrorBudget}
			if !errors.Is(waitErr, edgeservice.ErrSingleRequestInternalToolBudget) || terminal != want || terminalCount != 1 {
				t.Fatalf("iteration=%d Wait=%v terminal=%+v count=%d, want budget/1", iteration, waitErr, terminal, terminalCount)
			}
			if providerCalls.Load() != 1 || node.toolCount.Load() != 0 || node.cleanupCount.Load() != 0 || executor.bridge.pendingCount() != 0 {
				t.Fatalf("iteration=%d provider=%d tool=%d cleanup=%d pending=%d", iteration, providerCalls.Load(), node.toolCount.Load(), node.cleanupCount.Load(), executor.bridge.pendingCount())
			}
		}
	})

	t.Run("independent stage timeout", func(t *testing.T) {
		var providerCalls atomic.Int32
		mockSvc := &mockService{submit: func(ctx context.Context, _ edgeservice.ProviderPoolDispatchRequest) (*edgeservice.ProviderPoolDispatchResult, error) {
			providerCalls.Add(1)
			<-ctx.Done()
			return nil, ctx.Err()
		}}
		executor := NewSingleRequestExecutor(mockSvc)
		svc, binding, node := newTestServiceHarness(t, executor)
		binding.Limits.WallClockMS = 1000
		binding.Limits.StageTimeoutMS = 10

		execution, err := svc.StartSingleRequest(context.Background(), edgeservice.SingleRequestRequest{
			RequestID: "req-stage-timeout",
			Binding:   binding,
			Prompt:    "stage timeout ownership",
		})
		if err != nil {
			t.Fatal(err)
		}
		terminalCount := 0
		var terminal edgeservice.SingleRequestTerminalDisposition
		for progress := range execution.Progress() {
			if progress.Terminal != nil {
				terminalCount++
				terminal = *progress.Terminal
			}
		}
		_, waitErr := execution.Wait()
		want := edgeservice.SingleRequestTerminalDisposition{Kind: edgeservice.SingleRequestTerminalError, ErrorClass: edgeservice.SingleRequestTerminalErrorTimeout}
		if waitErr == nil || terminal != want || terminalCount != 1 {
			t.Fatalf("Wait=%v terminal=%+v count=%d, want timeout/1", waitErr, terminal, terminalCount)
		}
		if providerCalls.Load() != 1 || node.toolCount.Load() != 0 || node.cleanupCount.Load() != 0 || executor.bridge.pendingCount() != 0 {
			t.Fatalf("provider=%d tool=%d cleanup=%d pending=%d", providerCalls.Load(), node.toolCount.Load(), node.cleanupCount.Load(), executor.bridge.pendingCount())
		}
	})
}

func TestSingleRequestExecutorStageFailures(t *testing.T) {
	t.Run("PlanFailure", func(t *testing.T) {
		mockSvc := &mockService{
			submit: func(_ context.Context, _ edgeservice.ProviderPoolDispatchRequest) (*edgeservice.ProviderPoolDispatchResult, error) {
				return &edgeservice.ProviderPoolDispatchResult{
					Path:         edgeservice.ProviderPoolPathTunnel,
					Tunnel:       &mockTunnel{frames: framesFor(successBody("invalid plan json"))},
					DispatchInfo: matchingDispatch(),
				}, nil
			},
		}

		executor := NewSingleRequestExecutor(mockSvc)
		svc, binding, _ := newTestServiceHarness(t, executor)
		req := edgeservice.SingleRequestRequest{RequestID: "req-plan-fail", Binding: binding, Prompt: "Plan fail"}

		exec, err := svc.StartSingleRequest(context.Background(), req)
		if err != nil {
			t.Fatal(err)
		}

		_, err = waitExecutionResult(exec)
		if err == nil {
			t.Fatal("expected plan stage error, got nil")
		}
	})

	t.Run("WorkFailure", func(t *testing.T) {
		var callCount atomic.Int32
		mockSvc := &mockService{
			submit: func(_ context.Context, _ edgeservice.ProviderPoolDispatchRequest) (*edgeservice.ProviderPoolDispatchResult, error) {
				count := callCount.Add(1)
				if count == 1 {
					return &edgeservice.ProviderPoolDispatchResult{
						Path:         edgeservice.ProviderPoolPathTunnel,
						Tunnel:       &mockTunnel{frames: framesFor(executorPlanBody("Step 1", "Verify 1"))},
						DispatchInfo: matchingDispatch(),
					}, nil
				}
				return &edgeservice.ProviderPoolDispatchResult{
					Path:         edgeservice.ProviderPoolPathTunnel,
					Tunnel:       &mockTunnel{frames: framesFor(successBody("invalid work json"))},
					DispatchInfo: matchingDispatch(),
				}, nil
			},
		}

		executor := NewSingleRequestExecutor(mockSvc)
		svc, binding, _ := newTestServiceHarness(t, executor)
		req := edgeservice.SingleRequestRequest{RequestID: "req-work-fail", Binding: binding, Prompt: "Work fail"}

		exec, err := svc.StartSingleRequest(context.Background(), req)
		if err != nil {
			t.Fatal(err)
		}

		_, err = waitExecutionResult(exec)
		if err == nil {
			t.Fatal("expected work stage error, got nil")
		}
	})

	t.Run("ReviewFailure", func(t *testing.T) {
		var callCount atomic.Int32
		mockSvc := &mockService{
			submit: func(_ context.Context, _ edgeservice.ProviderPoolDispatchRequest) (*edgeservice.ProviderPoolDispatchResult, error) {
				count := callCount.Add(1)
				if count == 1 {
					return &edgeservice.ProviderPoolDispatchResult{
						Path:         edgeservice.ProviderPoolPathTunnel,
						Tunnel:       &mockTunnel{frames: framesFor(executorPlanBody("Step 1", "Verify 1"))},
						DispatchInfo: matchingDispatch(),
					}, nil
				}
				if count == 2 {
					return &edgeservice.ProviderPoolDispatchResult{
						Path:         edgeservice.ProviderPoolPathTunnel,
						Tunnel:       &mockTunnel{frames: framesFor(workToolBody("review-failure-work-1", edgeservice.InternalWorkspaceToolRead, `{"relative_path":"output.txt"}`))},
						DispatchInfo: matchingDispatch(),
					}, nil
				}
				if count == 3 {
					return &edgeservice.ProviderPoolDispatchResult{
						Path:         edgeservice.ProviderPoolPathTunnel,
						Tunnel:       &mockTunnel{frames: framesFor(executorWorkBody("Work done", "Work verified"))},
						DispatchInfo: matchingDispatch(),
					}, nil
				}
				return &edgeservice.ProviderPoolDispatchResult{
					Path:         edgeservice.ProviderPoolPathTunnel,
					Tunnel:       &mockTunnel{frames: framesFor(successBody("invalid review json"))},
					DispatchInfo: matchingDispatch(),
				}, nil
			},
		}

		executor := NewSingleRequestExecutor(mockSvc)
		svc, binding, _ := newTestServiceHarness(t, executor)
		req := edgeservice.SingleRequestRequest{RequestID: "req-review-fail", Binding: binding, Prompt: "Review fail"}

		exec, err := svc.StartSingleRequest(context.Background(), req)
		if err != nil {
			t.Fatal(err)
		}

		_, err = waitExecutionResult(exec)
		if err == nil {
			t.Fatal("expected review stage error, got nil")
		}
	})
}

func TestSingleRequestExecutorFinalOutputProvenance(t *testing.T) {
	responses := [][]byte{
		executorPlanBody("Plan step", "Plan verify"),
		workToolBody("work-provenance-1", edgeservice.InternalWorkspaceToolRead, `{"relative_path":"output.txt"}`),
		executorWorkBody("UNAPPROVED WORK CANDIDATE OUTPUT", "Work verified"),
		executorReviewPassBody("REVIEWER APPROVED TERMINAL OUTPUT", "Review approved"),
	}

	var callIndex atomic.Int32
	mockSvc := &mockService{
		submit: func(_ context.Context, _ edgeservice.ProviderPoolDispatchRequest) (*edgeservice.ProviderPoolDispatchResult, error) {
			idx := int(callIndex.Add(1) - 1)
			return &edgeservice.ProviderPoolDispatchResult{
				Path:         edgeservice.ProviderPoolPathTunnel,
				Tunnel:       &mockTunnel{frames: framesFor(responses[idx])},
				DispatchInfo: matchingDispatch(),
			}, nil
		},
	}

	executor := NewSingleRequestExecutor(mockSvc)
	svc, binding, _ := newTestServiceHarness(t, executor)
	req := edgeservice.SingleRequestRequest{RequestID: "req-provenance", Binding: binding, Prompt: "Provenance test"}

	exec, err := svc.StartSingleRequest(context.Background(), req)
	if err != nil {
		t.Fatal(err)
	}

	result, err := waitExecutionResult(exec)
	if err != nil {
		t.Fatal(err)
	}

	if result.Output != "REVIEWER APPROVED TERMINAL OUTPUT" {
		t.Fatalf("output = %q, want reviewer approved terminal output", result.Output)
	}
}

func TestSingleRequestExecutorTerminalWaiterCleanup(t *testing.T) {
	t.Run("Success", func(t *testing.T) {
		mockSvc := &mockService{
			submit: func(_ context.Context, req edgeservice.ProviderPoolDispatchRequest) (*edgeservice.ProviderPoolDispatchResult, error) {
				reqBody, _ := req.Tunnel.BuildBody("gemini-3.6-flash")
				bodyStr := string(reqBody)

				var resp []byte
				if strings.Contains(bodyStr, "Produce exactly one JSON object with non-empty string fields plan") {
					resp = executorPlanBody("Plan step", "Verify plan")
				} else if strings.Contains(bodyStr, "Read the supplied plan") {
					if !strings.Contains(bodyStr, "colliding-tool-id") {
						resp = workToolBody("colliding-tool-id", edgeservice.InternalWorkspaceToolRead, `{"relative_path":"output.txt"}`)
					} else {
						resp = executorWorkBody("Work done", "Work verified")
					}
				} else {
					resp = executorReviewPassBody("Success Output", "Approved")
				}

				return &edgeservice.ProviderPoolDispatchResult{
					Path:         edgeservice.ProviderPoolPathTunnel,
					Tunnel:       &mockTunnel{frames: framesFor(resp)},
					DispatchInfo: matchingDispatch(),
				}, nil
			},
		}

		executor := NewSingleRequestExecutor(mockSvc)
		svc, binding, _ := newTestServiceHarness(t, executor)

		req := edgeservice.SingleRequestRequest{
			RequestID: "req-clean-success",
			Binding:   binding,
			Prompt:    "Success task for req-clean-success",
		}

		exec, err := svc.StartSingleRequest(context.Background(), req)
		if err != nil {
			t.Fatal(err)
		}

		result, err := waitExecutionResult(exec)
		if err != nil {
			t.Fatalf("unexpected error: %v", err)
		}
		if result.Output != "Success Output" {
			t.Fatalf("got output %q, want %q", result.Output, "Success Output")
		}

		if executor.bridge.pendingCount() != 0 {
			t.Fatalf("pending count = %d, want 0 after successful completion", executor.bridge.pendingCount())
		}
	})

	t.Run("FailurePostRegistration", func(t *testing.T) {
		mockSvc := &mockService{
			submit: func(_ context.Context, req edgeservice.ProviderPoolDispatchRequest) (*edgeservice.ProviderPoolDispatchResult, error) {
				reqBody, _ := req.Tunnel.BuildBody("gemini-3.6-flash")
				bodyStr := string(reqBody)

				var resp []byte
				if strings.Contains(bodyStr, "Produce exactly one JSON object with non-empty string fields plan") {
					resp = executorPlanBody("Plan step", "Verify plan")
				} else if strings.Contains(bodyStr, "Read the supplied plan") {
					if !strings.Contains(bodyStr, "colliding-tool-id") {
						resp = workToolBody("colliding-tool-id", edgeservice.InternalWorkspaceToolRead, `{"relative_path":"output.txt"}`)
					} else {
						// Post-registration tool response: return invalid json to trigger Work stage failure
						resp = successBody("invalid work completion json")
					}
				} else {
					resp = executorReviewPassBody("Failure Output", "Approved")
				}

				return &edgeservice.ProviderPoolDispatchResult{
					Path:         edgeservice.ProviderPoolPathTunnel,
					Tunnel:       &mockTunnel{frames: framesFor(resp)},
					DispatchInfo: matchingDispatch(),
				}, nil
			},
		}

		executor := NewSingleRequestExecutor(mockSvc)
		svc, binding, _ := newTestServiceHarness(t, executor)

		req := edgeservice.SingleRequestRequest{
			RequestID: "req-clean-failure",
			Binding:   binding,
			Prompt:    "Failure task for req-clean-failure",
		}

		exec, err := svc.StartSingleRequest(context.Background(), req)
		if err != nil {
			t.Fatal(err)
		}

		_, err = waitExecutionResult(exec)
		if err == nil {
			t.Fatal("expected work stage failure, got nil")
		}

		if executor.bridge.pendingCount() != 0 {
			t.Fatalf("pending count = %d, want 0 after stage failure", executor.bridge.pendingCount())
		}
	})

	t.Run("CancellationWithActivePeer", func(t *testing.T) {
		cancelWaiterRegistered := make(chan struct{}, 1)
		unblockCancelTool := make(chan struct{})

		mockSvc := &mockService{
			submit: func(_ context.Context, req edgeservice.ProviderPoolDispatchRequest) (*edgeservice.ProviderPoolDispatchResult, error) {
				reqBody, _ := req.Tunnel.BuildBody("gemini-3.6-flash")
				bodyStr := string(reqBody)

				var reqID string
				if strings.Contains(bodyStr, "req-cancel-peer") {
					reqID = "req-cancel-peer"
				} else if strings.Contains(bodyStr, "req-active-peer") {
					reqID = "req-active-peer"
				}

				var resp []byte
				if strings.Contains(bodyStr, "Produce exactly one JSON object with non-empty string fields plan") {
					resp = executorPlanBody("Plan step for "+reqID, "Verify plan")
				} else if strings.Contains(bodyStr, "Read the supplied plan") {
					if !strings.Contains(bodyStr, "colliding-tool-id") {
						resp = workToolBody("colliding-tool-id", edgeservice.InternalWorkspaceToolRead, `{"relative_path":"output.txt"}`)
					} else {
						resp = executorWorkBody("Work done for "+reqID, "Work verified")
					}
				} else {
					resp = executorReviewPassBody("Approved for "+reqID, "Review approved")
				}

				return &edgeservice.ProviderPoolDispatchResult{
					Path:         edgeservice.ProviderPoolPathTunnel,
					Tunnel:       &mockTunnel{frames: framesFor(resp)},
					DispatchInfo: matchingDispatch(),
				}, nil
			},
		}

		executor := NewSingleRequestExecutor(mockSvc)
		svc, binding, nodeHarness := newTestServiceHarness(t, executor)

		var once sync.Once
		nodeHarness.toolResponder = func(req *iop.WorkspaceToolRequest) *iop.WorkspaceToolResponse {
			if req.GetRequestId() == "req-cancel-peer" {
				once.Do(func() {
					cancelWaiterRegistered <- struct{}{}
				})
				<-unblockCancelTool
			}
			return &iop.WorkspaceToolResponse{
				RequestId:  req.GetRequestId(),
				StageId:    req.GetStageId(),
				ToolCallId: req.GetToolCallId(),
				Status:     iop.WorkspaceStatus_WORKSPACE_STATUS_SUCCESS,
			}
		}

		ctxCancel, cancelFunc := context.WithCancel(context.Background())

		reqCancel := edgeservice.SingleRequestRequest{
			RequestID: "req-cancel-peer",
			Binding:   binding,
			Prompt:    "Task prompt for req-cancel-peer",
		}
		execCancel, err := svc.StartSingleRequest(ctxCancel, reqCancel)
		if err != nil {
			t.Fatal(err)
		}

		// Wait until req-cancel-peer has registered its waiter and reached toolResponder
		<-cancelWaiterRegistered

		// Cancel req-cancel-peer while its waiter is pending
		cancelFunc()

		// Start req-active-peer which reuses colliding-tool-id
		reqActive := edgeservice.SingleRequestRequest{
			RequestID: "req-active-peer",
			Binding:   binding,
			Prompt:    "Task prompt for req-active-peer",
		}
		execActive, err := svc.StartSingleRequest(context.Background(), reqActive)
		if err != nil {
			t.Fatal(err)
		}

		// Active peer should complete successfully
		resultActive, errActive := waitExecutionResult(execActive)
		if errActive != nil {
			t.Fatalf("active peer error: %v", errActive)
		}
		if resultActive.Output != "Approved for req-active-peer" {
			t.Fatalf("active peer output = %q, want %q", resultActive.Output, "Approved for req-active-peer")
		}

		// Unblock cancelled tool responder so goroutine finishes
		close(unblockCancelTool)

		_, errCancel := waitExecutionResult(execCancel)
		if !errors.Is(errCancel, edgeservice.ErrSingleRequestCancelled) && !errors.Is(errCancel, context.Canceled) {
			t.Fatalf("expected cancel error for req-cancel-peer, got: %v", errCancel)
		}

		// Success, failure, and cancellation run through StartSingleRequest with a
		// registered continuation waiter; no test calls clearRequest directly.
		if executor.bridge.pendingCount() != 0 {
			t.Fatalf("pending count = %d, want 0 after cancellation with active peer", executor.bridge.pendingCount())
		}
	})
}