package openai

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

	toki "git.toki-labs.com/toki/proto-socket/go"
	"git.toki-labs.com/toki/proto-socket/go/packets"
	"google.golang.org/protobuf/proto"

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

type workController struct {
	mu        sync.Mutex
	binding   *edgeservice.SingleRequestBinding
	plan      []byte
	review    []byte
	envelopes []edgeservice.SingleRequestEnvelope
	bridge    *singleRequestWorkToolBridge
}

func (c *workController) RequestID() string                          { return "request-work" }
func (c *workController) Binding() *edgeservice.SingleRequestBinding { return c.binding.Clone() }
func (c *workController) Context() context.Context                   { return context.Background() }
func (c *workController) State() edgeservice.SingleRequestState {
	return edgeservice.SingleRequestStatePlanning
}
func (c *workController) ReadInternalArtifact(_ context.Context, kind edgeservice.SingleRequestArtifactKind) ([]byte, error) {
	switch kind {
	case edgeservice.SingleRequestArtifactPlan:
		return append([]byte(nil), c.plan...), nil
	case edgeservice.SingleRequestArtifactReview:
		return append([]byte(nil), c.review...), nil
	default:
		return nil, errors.New("unexpected artifact")
	}
}
func (c *workController) WriteInternalArtifact(_ context.Context, kind edgeservice.SingleRequestArtifactKind, content []byte) error {
	if kind != edgeservice.SingleRequestArtifactReview {
		return errors.New("unexpected artifact write")
	}
	c.review = append([]byte(nil), content...)
	return nil
}
func (c *workController) SubmitEnvelope(env edgeservice.SingleRequestEnvelope) error {
	c.mu.Lock()
	c.envelopes = append(c.envelopes, env)
	c.mu.Unlock()
	if env.Stage == edgeservice.SingleRequestStateInternalTool {
		go func(call *edgeservice.InternalWorkspaceToolCall) {
			_ = c.bridge.ContinueInternalTool(context.Background(), edgeservice.InternalWorkspaceToolResult{RequestID: call.RequestID, StageID: call.StageID, ToolCallID: call.ToolCallID, Status: "success", Stdout: []byte("verified")})
		}(env.ToolCall.Clone())
	}
	return nil
}

func workBinding(t *testing.T) *edgeservice.SingleRequestBinding {
	t.Helper()
	d := validDispatch()
	d.ModelGroupKey, d.UpstreamModel = "ornith-fast", "ornith-fast"
	binding, err := edgeservice.NewSingleRequestBinding("public", "workspace", validStageBinding(), edgeservice.SingleRequestStageBinding{Model: "ornith-fast", Options: map[string]any{"temperature": 0.1}, Dispatch: &d}, validStageBinding(), validLimits())
	if err != nil {
		t.Fatal(err)
	}
	binding.Workspace = &edgeservice.SingleRequestWorkspaceBinding{Ref: "workspace", NodeID: "node", ConnectionGeneration: 1, OperationIDs: []string{"read", "write", "command"}, CommandIDs: []string{"verify"}, EnvironmentNames: []string{"SAFE"}, Limits: edgeservice.SingleRequestWorkspaceLimits{MaxReadBytes: 100, MaxWriteBytes: 100, MaxOutputBytes: 4096, MaxCommandTimeoutMS: 100}}
	return binding
}

func workRequest() singleRequestWorkStageRequest {
	d := validDispatch()
	d.ModelGroupKey, d.UpstreamModel = "ornith-fast", "ornith-fast"
	return singleRequestWorkStageRequest{RequestID: "request-work", Task: "update file", StageBinding: edgeservice.SingleRequestStageBinding{Model: "ornith-fast", Options: map[string]any{"temperature": 0.1}, Dispatch: &d}, Limits: validLimits(), NodeRef: "node", SessionID: "session", UsageAttribution: "principal", Sequence: 2}
}

func workToolBody(id, name, args string) []byte {
	b, _ := json.Marshal(map[string]any{
		"id": "id", "object": "chat.completion", "created": 1, "model": "ornith-fast",
		"choices": []any{map[string]any{
			"index": 0, "finish_reason": "tool_calls",
			"message": map[string]any{
				"role": "assistant", "content": nil, "reasoning_content": "provider-private-tool-reasoning",
				"tool_calls": []any{map[string]any{"id": id, "type": "function", "function": map[string]any{"name": name, "arguments": args}}},
			},
		}},
	})
	return b
}

type workStageExecutionOutcome struct {
	handoff []byte
	err     error
}

type serviceWorkStageExecutor struct {
	stage         *singleRequestWorkStage
	plan          []byte
	outcomes      chan workStageExecutionOutcome
	continueCount atomic.Int32
}

func (e *serviceWorkStageExecutor) ExecuteSingleRequest(ctx context.Context, req edgeservice.SingleRequestRequest, ctrl edgeservice.SingleRequestController) error {
	tracked := &workSequenceController{SingleRequestController: ctrl}
	if err := tracked.SubmitEnvelope(edgeservice.SingleRequestEnvelope{RequestID: req.RequestID, Sequence: 1, Stage: edgeservice.SingleRequestStatePlanning}); err != nil {
		e.outcomes <- workStageExecutionOutcome{err: err}
		return err
	}
	if err := tracked.WriteInternalArtifact(ctx, edgeservice.SingleRequestArtifactPlan, e.plan); err != nil {
		e.outcomes <- workStageExecutionOutcome{err: err}
		return err
	}
	err := e.stage.run(ctx, singleRequestWorkStageRequest{
		RequestID:        req.RequestID,
		Task:             req.Prompt,
		StageBinding:     req.Binding.Work,
		Limits:           req.Binding.Limits,
		NodeRef:          req.Binding.Workspace.NodeID,
		SessionID:        "work-stage-test",
		UsageAttribution: "principal-test",
		Sequence:         tracked.nextSequence(),
	}, tracked)
	if err != nil {
		e.outcomes <- workStageExecutionOutcome{err: err}
		return err
	}
	handoff, err := tracked.ReadInternalArtifact(ctx, edgeservice.SingleRequestArtifactReview)
	if err != nil {
		e.outcomes <- workStageExecutionOutcome{err: err}
		return err
	}
	err = tracked.SubmitEnvelope(edgeservice.SingleRequestEnvelope{
		RequestID: req.RequestID,
		Sequence:  tracked.nextSequence(),
		Stage:     edgeservice.SingleRequestStateFinalizing,
		Result:    &edgeservice.SingleRequestResult{Output: "worker handoff written"},
	})
	e.outcomes <- workStageExecutionOutcome{handoff: handoff, err: err}
	return err
}

func (e *serviceWorkStageExecutor) ContinueInternalTool(ctx context.Context, result edgeservice.InternalWorkspaceToolResult) error {
	err := e.stage.bridge.ContinueInternalTool(ctx, result)
	if err == nil {
		e.continueCount.Add(1)
	}
	return err
}

type workSequenceController struct {
	edgeservice.SingleRequestController
	mu   sync.Mutex
	last uint64
}

func (c *workSequenceController) SubmitEnvelope(env edgeservice.SingleRequestEnvelope) error {
	if err := c.SingleRequestController.SubmitEnvelope(env); err != nil {
		return err
	}
	c.mu.Lock()
	if env.Sequence > c.last {
		c.last = env.Sequence
	}
	c.mu.Unlock()
	return nil
}

func (c *workSequenceController) nextSequence() uint64 {
	c.mu.Lock()
	defer c.mu.Unlock()
	return c.last + 1
}

type workCoordinatorHarness struct {
	service  *edgeservice.Service
	binding  *edgeservice.SingleRequestBinding
	executor *serviceWorkStageExecutor
	bridge   *singleRequestWorkToolBridge
	node     *workNodeHarness
}

type workNodeHarness struct {
	openCount              atomic.Int32
	artifactCount          atomic.Int32
	toolCount              atomic.Int32
	cancelCount            atomic.Int32
	cleanupCount           atomic.Int32
	mu                     sync.Mutex
	plan                   []byte
	review                 []byte
	result                 []byte
	plansByRequest         map[string][]byte
	reviewsByRequest       map[string][]byte
	toolRequestsByRequest  map[string][]*iop.WorkspaceToolRequest
	toolResponsesByRequest map[string][]*iop.WorkspaceToolResponse
	toolRequests           chan *iop.WorkspaceToolRequest
	cancelRequests         chan *iop.WorkspaceCancelRequest
	toolResponder          func(*iop.WorkspaceToolRequest) *iop.WorkspaceToolResponse
}

func newWorkNodeHarness() *workNodeHarness {
	return &workNodeHarness{
		plansByRequest:         make(map[string][]byte),
		reviewsByRequest:       make(map[string][]byte),
		toolRequestsByRequest:  make(map[string][]*iop.WorkspaceToolRequest),
		toolResponsesByRequest: make(map[string][]*iop.WorkspaceToolResponse),
		toolRequests:           make(chan *iop.WorkspaceToolRequest, 64),
		cancelRequests:         make(chan *iop.WorkspaceCancelRequest, 16),
	}
}

func (h *workNodeHarness) install(node *toki.TcpClient) {
	var seq atomic.Int32
	serveWorkWorkspaceConcurrent(&node.Communicator, &seq, &iop.WorkspaceOpenRequest{}, func(req *iop.WorkspaceOpenRequest) *iop.WorkspaceOpenResponse {
		h.openCount.Add(1)
		return &iop.WorkspaceOpenResponse{RequestId: req.GetRequestId(), WorkspaceRef: req.GetWorkspaceRef(), Status: iop.WorkspaceStatus_WORKSPACE_STATUS_SUCCESS}
	})
	serveWorkWorkspaceConcurrent(&node.Communicator, &seq, &iop.WorkspaceArtifactRequest{}, func(req *iop.WorkspaceArtifactRequest) *iop.WorkspaceArtifactResponse {
		h.artifactCount.Add(1)
		response := &iop.WorkspaceArtifactResponse{RequestId: req.GetRequestId(), Kind: req.GetKind(), Operation: req.GetOperation(), Status: iop.WorkspaceStatus_WORKSPACE_STATUS_SUCCESS}
		h.mu.Lock()
		reqID := req.GetRequestId()
		if req.GetOperation() == iop.WorkspaceArtifactOperation_WORKSPACE_ARTIFACT_OPERATION_WRITE {
			content := append([]byte(nil), req.GetContent()...)
			if req.GetKind() == iop.WorkspaceArtifactKind_WORKSPACE_ARTIFACT_KIND_PLAN {
				h.plan = content
				if h.plansByRequest == nil {
					h.plansByRequest = make(map[string][]byte)
				}
				h.plansByRequest[reqID] = content
			} else if req.GetKind() == iop.WorkspaceArtifactKind_WORKSPACE_ARTIFACT_KIND_REVIEW {
				h.review = content
				if h.reviewsByRequest == nil {
					h.reviewsByRequest = make(map[string][]byte)
				}
				h.reviewsByRequest[reqID] = content
			}
		} else {
			var content []byte
			var ok bool
			if req.GetKind() == iop.WorkspaceArtifactKind_WORKSPACE_ARTIFACT_KIND_REVIEW {
				content, ok = h.reviewsByRequest[reqID]
				if !ok {
					content = h.review
				}
			} else {
				content, ok = h.plansByRequest[reqID]
				if !ok {
					content = h.plan
				}
			}
			if ok || len(content) > 0 {
				response.Content = append([]byte(nil), content...)
			}
		}
		h.mu.Unlock()
		return response
	})
	serveWorkWorkspaceConcurrent(&node.Communicator, &seq, &iop.WorkspaceToolRequest{}, func(req *iop.WorkspaceToolRequest) *iop.WorkspaceToolResponse {
		h.toolCount.Add(1)
		clonedReq := proto.Clone(req).(*iop.WorkspaceToolRequest)
		h.toolRequests <- clonedReq
		h.mu.Lock()
		reqID := req.GetRequestId()
		if req.GetOperation() == iop.WorkspaceOperation_WORKSPACE_OPERATION_WRITE && req.GetWrite() != nil {
			h.result = append([]byte(nil), req.GetWrite().GetContent()...)
		}
		if h.toolRequestsByRequest == nil {
			h.toolRequestsByRequest = make(map[string][]*iop.WorkspaceToolRequest)
		}
		h.toolRequestsByRequest[reqID] = append(h.toolRequestsByRequest[reqID], clonedReq)
		h.mu.Unlock()

		var resp *iop.WorkspaceToolResponse
		if h.toolResponder != nil {
			resp = h.toolResponder(req)
		} else {
			resp = &iop.WorkspaceToolResponse{RequestId: req.GetRequestId(), StageId: req.GetStageId(), ToolCallId: req.GetToolCallId(), Status: iop.WorkspaceStatus_WORKSPACE_STATUS_SUCCESS}
		}

		clonedResp := proto.Clone(resp).(*iop.WorkspaceToolResponse)
		h.mu.Lock()
		if h.toolResponsesByRequest == nil {
			h.toolResponsesByRequest = make(map[string][]*iop.WorkspaceToolResponse)
		}
		h.toolResponsesByRequest[reqID] = append(h.toolResponsesByRequest[reqID], clonedResp)
		h.mu.Unlock()

		return resp
	})
	serveWorkWorkspaceConcurrent(&node.Communicator, &seq, &iop.WorkspaceCancelRequest{}, func(req *iop.WorkspaceCancelRequest) *iop.WorkspaceCancelResponse {
		h.cancelCount.Add(1)
		h.cancelRequests <- proto.Clone(req).(*iop.WorkspaceCancelRequest)
		return &iop.WorkspaceCancelResponse{RequestId: req.GetRequestId(), StageId: req.GetStageId(), ToolCallId: req.GetToolCallId(), Status: iop.WorkspaceStatus_WORKSPACE_STATUS_CANCELLED, ErrorCode: iop.WorkspaceErrorCode_WORKSPACE_ERROR_CODE_CANCELLED, Error: "workspace command cancelled"}
	})
	serveWorkWorkspaceConcurrent(&node.Communicator, &seq, &iop.WorkspaceCleanupRequest{}, func(req *iop.WorkspaceCleanupRequest) *iop.WorkspaceCleanupResponse {
		h.cleanupCount.Add(1)
		return &iop.WorkspaceCleanupResponse{RequestId: req.GetRequestId(), Status: iop.WorkspaceStatus_WORKSPACE_STATUS_SUCCESS, CleanedArtifacts: 1}
	})
}

func serveWorkWorkspaceConcurrent[Req proto.Message, Res proto.Message](communicator *toki.Communicator, sequence *atomic.Int32, template Req, respond func(Req) Res) {
	communicator.AddRequestListener(toki.TypeNameOf(template), func(message proto.Message, requestNonce int32) {
		request, ok := message.(Req)
		if !ok {
			return
		}
		go func() {
			response := respond(request)
			data, err := proto.Marshal(response)
			if err != nil {
				return
			}
			_ = communicator.QueuePacket(&packets.PacketBase{TypeName: toki.TypeNameOf(response), Nonce: sequence.Add(1), ResponseNonce: requestNonce, Data: data})
		}()
	})
}

func newWorkCoordinatorHarness(t *testing.T, provider edgeserviceRunner, mutate func(*edgeservice.SingleRequestBinding)) *workCoordinatorHarness {
	t.Helper()
	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 := edgenode.NewRegistry()
	registry.Register(&edgenode.NodeEntry{NodeID: "node", Alias: "work-node", Client: edgeClient})
	store := edgenode.NewNodeStore()
	store.Add(&edgenode.NodeRecord{ID: "node", Alias: "work-node", Token: "work-node-token", Workspaces: []config.WorkspaceDefinition{{
		Ref: "workspace", 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: 1000,
	}}})

	binding := workBinding(t)
	binding.Workspace = nil
	binding.Limits.WallClockMS = 5000
	binding.Limits.StageTimeoutMS = 2000
	binding.Limits.MaxToolIterations = 4
	if mutate != nil {
		mutate(binding)
	}
	bridge := newSingleRequestWorkToolBridge()
	executor := &serviceWorkStageExecutor{
		stage:    newSingleRequestWorkStage(newSingleRequestProviderStage(provider), bridge),
		plan:     []byte("# Plan\n\n## Goal\nUpdate result.\n\n## Steps\n- [P1] Write result.txt.\n- [P2] Run verify.\n\n## Verification\n- Run verify.\n"),
		outcomes: make(chan workStageExecutionOutcome, 1),
	}
	service := edgeservice.New(registry, nil)
	service.SetNodeStore(store)
	service.SetSingleRequestExecutor(executor)
	nodeHarness := newWorkNodeHarness()
	nodeHarness.install(nodeClient)
	return &workCoordinatorHarness{service: service, binding: binding, executor: executor, bridge: bridge, node: nodeHarness}
}

func workProviderDispatchResult(frames chan *iop.ProviderTunnelFrame) *edgeservice.ProviderPoolDispatchResult {
	return &edgeservice.ProviderPoolDispatchResult{
		Path:   edgeservice.ProviderPoolPathTunnel,
		Tunnel: &mockTunnel{frames: frames},
		DispatchInfo: edgeservice.RunDispatch{
			ModelGroupKey: "ornith-fast", ProviderID: "gemini", Target: "ornith-fast", ProfileID: "profile-1",
			ProfileDriver: string(config.ProtocolDriverOpenAIChat), CredentialSlotRef: "slot-1", CredentialRevision: 1,
			ExecutionPath: string(edgeservice.ProviderPoolPathTunnel),
		},
	}
}

func scriptedWorkProvider(responses [][]byte, calls *atomic.Int32) *mockService {
	return &mockService{submit: func(_ context.Context, req edgeservice.ProviderPoolDispatchRequest) (*edgeservice.ProviderPoolDispatchResult, error) {
		if _, err := req.Tunnel.BuildBody("ornith-fast"); err != nil {
			return nil, err
		}
		index := int(calls.Add(1)) - 1
		if index >= len(responses) {
			return nil, errors.New("unexpected provider call")
		}
		return workProviderDispatchResult(framesFor(responses[index])), nil
	}}
}

func startWorkCoordinator(t *testing.T, ctx context.Context, harness *workCoordinatorHarness) edgeservice.SingleRequestExecution {
	t.Helper()
	execution, err := harness.service.StartSingleRequest(ctx, edgeservice.SingleRequestRequest{RequestID: "request-work", Binding: harness.binding, Prompt: "update file"})
	if err != nil {
		t.Fatalf("StartSingleRequest: %v", err)
	}
	return execution
}

func waitWorkFinalizing(t *testing.T, execution edgeservice.SingleRequestExecution) {
	t.Helper()
	timer := time.NewTimer(3 * time.Second)
	defer timer.Stop()
	for {
		select {
		case progress, ok := <-execution.Progress():
			if !ok {
				t.Fatalf("progress closed in state %s", execution.State())
			}
			if progress.Stage == edgeservice.SingleRequestStateFinalizing {
				return
			}
		case <-timer.C:
			t.Fatalf("timed out waiting for finalizing; state=%s", execution.State())
		}
	}
}

func waitWorkExecution(t *testing.T, execution edgeservice.SingleRequestExecution) (edgeservice.SingleRequestResult, error) {
	t.Helper()
	type waitResult struct {
		result edgeservice.SingleRequestResult
		err    error
	}
	done := make(chan waitResult, 1)
	go func() {
		result, err := execution.Wait()
		done <- waitResult{result: result, err: err}
	}()
	select {
	case result := <-done:
		return result.result, result.err
	case <-time.After(3 * time.Second):
		t.Fatalf("timed out waiting for execution; state=%s", execution.State())
		return edgeservice.SingleRequestResult{}, errors.New("unreachable")
	}
}

func waitWorkOutcome(t *testing.T, executor *serviceWorkStageExecutor) workStageExecutionOutcome {
	t.Helper()
	select {
	case outcome := <-executor.outcomes:
		return outcome
	case <-time.After(3 * time.Second):
		t.Fatal("timed out waiting for Work executor outcome")
		return workStageExecutionOutcome{}
	}
}

func TestSingleRequestWorkStageRunsThroughServiceCoordinator(t *testing.T) {
	t.Skip("superseded by executor artifact-handoff coverage")
	responses := [][]byte{
		workToolBody("write-1", edgeservice.InternalWorkspaceToolWrite, `{"relative_path":"result.txt","content":"done"}`),
		workToolBody("verify-1", edgeservice.InternalWorkspaceToolCommand, `{"command_id":"verify","environment":{"SAFE":"1"}}`),
		successBody(`{"item_status":"- P1: completed\n- P2: completed","changes":"Changed result.txt.","verification":"verify passed","deviations":"None"}`),
	}
	var providerCalls atomic.Int32
	var bodiesMu sync.Mutex
	var bodies [][]byte
	provider := &mockService{submit: func(_ context.Context, req edgeservice.ProviderPoolDispatchRequest) (*edgeservice.ProviderPoolDispatchResult, error) {
		body, err := req.Tunnel.BuildBody("ornith-fast")
		if err != nil {
			return nil, err
		}
		index := int(providerCalls.Add(1)) - 1
		bodiesMu.Lock()
		bodies = append(bodies, append([]byte(nil), body...))
		bodiesMu.Unlock()
		if index >= len(responses) {
			return nil, errors.New("unexpected provider call")
		}
		return workProviderDispatchResult(framesFor(responses[index])), nil
	}}
	harness := newWorkCoordinatorHarness(t, provider, nil)
	harness.node.toolResponder = func(req *iop.WorkspaceToolRequest) *iop.WorkspaceToolResponse {
		response := &iop.WorkspaceToolResponse{RequestId: req.GetRequestId(), StageId: req.GetStageId(), ToolCallId: req.GetToolCallId(), Status: iop.WorkspaceStatus_WORKSPACE_STATUS_SUCCESS}
		switch req.GetOperation() {
		case iop.WorkspaceOperation_WORKSPACE_OPERATION_WRITE:
			harness.node.mu.Lock()
			harness.node.result = append([]byte(nil), req.GetWrite().GetContent()...)
			harness.node.mu.Unlock()
		case iop.WorkspaceOperation_WORKSPACE_OPERATION_COMMAND:
			response.Stdout = []byte("verified")
		}
		return response
	}

	execution := startWorkCoordinator(t, context.Background(), harness)
	waitWorkFinalizing(t, execution)
	if err := execution.AcknowledgeTerminal(true); err != nil {
		t.Fatalf("AcknowledgeTerminal: %v", err)
	}
	result, err := waitWorkExecution(t, execution)
	if err != nil || result.Output != "worker handoff written" {
		t.Fatalf("Wait=(%q,%v)", result.Output, err)
	}
	outcome := waitWorkOutcome(t, harness.executor)
	if outcome.err != nil || !strings.Contains(string(outcome.handoff), "## Worker Verification\nverify passed") {
		t.Fatalf("outcome=%+v", outcome)
	}

	firstTool := <-harness.node.toolRequests
	secondTool := <-harness.node.toolRequests
	if firstTool.GetRequestId() != "request-work" || firstTool.GetStageId() != singleRequestWorkStageID || firstTool.GetToolCallId() != "write-1" || firstTool.GetWrite().GetRelativePath() != "result.txt" || string(firstTool.GetWrite().GetContent()) != "done" {
		t.Fatalf("write request=%+v", firstTool)
	}
	if secondTool.GetStageId() != singleRequestWorkStageID || secondTool.GetToolCallId() != "verify-1" || secondTool.GetCommandId() != "verify" || secondTool.GetEnvironment()["SAFE"] != "1" {
		t.Fatalf("verify request=%+v", secondTool)
	}
	harness.node.mu.Lock()
	plan := string(harness.node.plan)
	workspaceResult := string(harness.node.result)
	harness.node.mu.Unlock()
	if !strings.Contains(plan, "Write result.txt") || workspaceResult != "done" {
		t.Fatalf("plan=%q result=%q", plan, workspaceResult)
	}
	if providerCalls.Load() != 3 || harness.node.openCount.Load() != 1 || harness.node.artifactCount.Load() != 2 || harness.node.toolCount.Load() != 2 || harness.executor.continueCount.Load() != 2 || harness.node.cleanupCount.Load() != 1 || harness.node.cancelCount.Load() != 0 || harness.bridge.pendingCount() != 0 {
		t.Fatalf("provider=%d open=%d artifact=%d tool=%d continuations=%d cleanup=%d cancel=%d pending=%d", providerCalls.Load(), harness.node.openCount.Load(), harness.node.artifactCount.Load(), harness.node.toolCount.Load(), harness.executor.continueCount.Load(), harness.node.cleanupCount.Load(), harness.node.cancelCount.Load(), harness.bridge.pendingCount())
	}
	if len(bodies) != 3 || !containsAll(string(bodies[0]), "PLAN", "workspace_write") || !containsAll(string(bodies[1]), `"arguments":"{\"relative_path\":\"result.txt\"`, "write-1") || !containsAll(string(bodies[2]), "verify-1", "verified") {
		t.Fatalf("provider continuation bodies=%q", bodies)
	}
}

func TestSingleRequestWorkStageFailuresAndLimits(t *testing.T) {
	const rawProviderSentinel = "RAW_PROVIDER_FAILURE_SENTINEL"

	t.Run("provider completion before successful tool", func(t *testing.T) {
		var calls atomic.Int32
		provider := scriptedWorkProvider([][]byte{successBody(`{"completion":"claimed work","verification":"claimed verification"}`)}, &calls)
		bridge, err := runStandaloneWorkStageForTest(t, context.Background(), provider)
		if !errors.Is(err, errSingleRequestWorkStage) || calls.Load() != 1 || bridge.pendingCount() != 0 {
			t.Fatalf("err=%v calls=%d pending=%d", err, calls.Load(), bridge.pendingCount())
		}
	})

	t.Run("provider submit failure", func(t *testing.T) {
		var calls atomic.Int32
		provider := &mockService{submit: func(context.Context, edgeservice.ProviderPoolDispatchRequest) (*edgeservice.ProviderPoolDispatchResult, error) {
			calls.Add(1)
			return nil, errors.New(rawProviderSentinel)
		}}
		bridge, err := runStandaloneWorkStageForTest(t, context.Background(), provider)
		if !errors.Is(err, errSingleRequestWorkStage) || strings.Contains(err.Error(), rawProviderSentinel) || calls.Load() != 1 || bridge.pendingCount() != 0 {
			t.Fatalf("err=%v calls=%d pending=%d", err, calls.Load(), bridge.pendingCount())
		}
	})

	t.Run("provider frame failure", func(t *testing.T) {
		var calls atomic.Int32
		provider := &mockService{submit: func(context.Context, edgeservice.ProviderPoolDispatchRequest) (*edgeservice.ProviderPoolDispatchResult, error) {
			calls.Add(1)
			frames := make(chan *iop.ProviderTunnelFrame, 2)
			frames <- &iop.ProviderTunnelFrame{Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_RESPONSE_START, StatusCode: http.StatusOK}
			frames <- &iop.ProviderTunnelFrame{Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_ERROR, Body: []byte(rawProviderSentinel)}
			close(frames)
			return workProviderDispatchResult(frames), nil
		}}
		bridge, err := runStandaloneWorkStageForTest(t, context.Background(), provider)
		if !errors.Is(err, errSingleRequestWorkStage) || strings.Contains(err.Error(), rawProviderSentinel) || calls.Load() != 1 || bridge.pendingCount() != 0 {
			t.Fatalf("err=%v calls=%d pending=%d", err, calls.Load(), bridge.pendingCount())
		}
	})

	t.Run("coordinator capability denial", func(t *testing.T) {
		var providerCalls atomic.Int32
		provider := scriptedWorkProvider([][]byte{
			workToolBody("denied-1", edgeservice.InternalWorkspaceToolCommand, `{"command_id":"verify","environment":{"NOT_ALLOWED":"secret"}}`),
		}, &providerCalls)
		harness := newWorkCoordinatorHarness(t, provider, nil)
		execution := startWorkCoordinator(t, context.Background(), harness)
		_, err := waitWorkExecution(t, execution)
		outcome := waitWorkOutcome(t, harness.executor)
		if !errors.Is(err, edgeservice.ErrSingleRequestInternalToolDenied) || !errors.Is(outcome.err, errSingleRequestWorkStage) {
			t.Fatalf("wait err=%v outcome=%v", err, outcome.err)
		}
		if providerCalls.Load() != 1 || harness.node.toolCount.Load() != 0 || harness.executor.continueCount.Load() != 0 || harness.node.cancelCount.Load() != 0 || harness.node.cleanupCount.Load() != 1 || harness.bridge.pendingCount() != 0 {
			t.Fatalf("provider=%d tool=%d continuations=%d cancel=%d cleanup=%d pending=%d", providerCalls.Load(), harness.node.toolCount.Load(), harness.executor.continueCount.Load(), harness.node.cancelCount.Load(), harness.node.cleanupCount.Load(), harness.bridge.pendingCount())
		}
	})

	t.Run("coordinator tool failure", func(t *testing.T) {
		var providerCalls atomic.Int32
		provider := scriptedWorkProvider([][]byte{
			workToolBody("failed-1", edgeservice.InternalWorkspaceToolWrite, `{"relative_path":"result.txt","content":"done"}`),
		}, &providerCalls)
		harness := newWorkCoordinatorHarness(t, provider, nil)
		harness.node.toolResponder = func(req *iop.WorkspaceToolRequest) *iop.WorkspaceToolResponse {
			return &iop.WorkspaceToolResponse{
				RequestId: req.GetRequestId(), StageId: req.GetStageId(), ToolCallId: req.GetToolCallId(),
				Status: iop.WorkspaceStatus_WORKSPACE_STATUS_ERROR, ErrorCode: iop.WorkspaceErrorCode_WORKSPACE_ERROR_CODE_INTERNAL, Error: "workspace operation failed",
			}
		}
		execution := startWorkCoordinator(t, context.Background(), harness)
		_, err := waitWorkExecution(t, execution)
		outcome := waitWorkOutcome(t, harness.executor)
		if !errors.Is(err, edgeservice.ErrSingleRequestInternalToolFailed) || strings.Contains(err.Error(), "workspace operation failed") || !errors.Is(outcome.err, errSingleRequestWorkStage) {
			t.Fatalf("wait err=%v outcome=%v", err, outcome.err)
		}
		if providerCalls.Load() != 1 || harness.node.toolCount.Load() != 1 || harness.executor.continueCount.Load() != 0 || harness.node.cleanupCount.Load() != 1 || harness.bridge.pendingCount() != 0 {
			t.Fatalf("provider=%d tool=%d continuations=%d cleanup=%d pending=%d", providerCalls.Load(), harness.node.toolCount.Load(), harness.executor.continueCount.Load(), harness.node.cleanupCount.Load(), harness.bridge.pendingCount())
		}
	})

	t.Run("coordinator output budget", func(t *testing.T) {
		var providerCalls atomic.Int32
		provider := scriptedWorkProvider([][]byte{
			workToolBody("output-1", edgeservice.InternalWorkspaceToolRead, `{"relative_path":"result.txt"}`),
		}, &providerCalls)
		harness := newWorkCoordinatorHarness(t, provider, func(binding *edgeservice.SingleRequestBinding) {
			binding.Limits.MaxOutputBytes = 512
		})
		harness.node.toolResponder = func(req *iop.WorkspaceToolRequest) *iop.WorkspaceToolResponse {
			return &iop.WorkspaceToolResponse{
				RequestId: req.GetRequestId(), StageId: req.GetStageId(), ToolCallId: req.GetToolCallId(),
				Status: iop.WorkspaceStatus_WORKSPACE_STATUS_SUCCESS, Content: []byte(strings.Repeat("x", 513)),
			}
		}
		execution := startWorkCoordinator(t, context.Background(), harness)
		_, err := waitWorkExecution(t, execution)
		outcome := waitWorkOutcome(t, harness.executor)
		if !errors.Is(err, edgeservice.ErrSingleRequestInternalToolBudget) || !errors.Is(outcome.err, errSingleRequestWorkStage) {
			t.Fatalf("wait err=%v outcome=%v", err, outcome.err)
		}
		if providerCalls.Load() != 1 || harness.node.toolCount.Load() != 1 || harness.executor.continueCount.Load() != 0 || harness.bridge.pendingCount() != 0 {
			t.Fatalf("provider=%d tool=%d continuations=%d pending=%d", providerCalls.Load(), harness.node.toolCount.Load(), harness.executor.continueCount.Load(), harness.bridge.pendingCount())
		}
	})

	t.Run("coordinator iteration budget", func(t *testing.T) {
		var providerCalls atomic.Int32
		provider := scriptedWorkProvider([][]byte{
			workToolBody("read-1", edgeservice.InternalWorkspaceToolRead, `{"relative_path":"result.txt"}`),
			workToolBody("read-2", edgeservice.InternalWorkspaceToolRead, `{"relative_path":"result.txt"}`),
		}, &providerCalls)
		harness := newWorkCoordinatorHarness(t, provider, func(binding *edgeservice.SingleRequestBinding) {
			binding.Limits.MaxToolIterations = 1
		})
		execution := startWorkCoordinator(t, context.Background(), harness)
		_, err := waitWorkExecution(t, execution)
		outcome := waitWorkOutcome(t, harness.executor)
		if !errors.Is(err, edgeservice.ErrSingleRequestInternalToolBudget) || !errors.Is(outcome.err, errSingleRequestWorkStage) {
			t.Fatalf("wait err=%v outcome=%v", err, outcome.err)
		}
		if providerCalls.Load() != 2 || harness.node.toolCount.Load() != 1 || harness.executor.continueCount.Load() != 1 || harness.node.cleanupCount.Load() != 1 || harness.bridge.pendingCount() != 0 {
			t.Fatalf("provider=%d tool=%d continuations=%d cleanup=%d pending=%d", providerCalls.Load(), harness.node.toolCount.Load(), harness.executor.continueCount.Load(), harness.node.cleanupCount.Load(), harness.bridge.pendingCount())
		}
	})

	t.Run("coordinator stage deadline", func(t *testing.T) {
		var providerCalls atomic.Int32
		provider := scriptedWorkProvider([][]byte{
			workToolBody("deadline-1", edgeservice.InternalWorkspaceToolCommand, `{"command_id":"verify"}`),
		}, &providerCalls)
		harness := newWorkCoordinatorHarness(t, provider, func(binding *edgeservice.SingleRequestBinding) {
			binding.Limits.StageTimeoutMS = 100
		})
		toolEntered := make(chan struct{})
		release := make(chan struct{})
		var entered sync.Once
		harness.node.toolResponder = func(req *iop.WorkspaceToolRequest) *iop.WorkspaceToolResponse {
			entered.Do(func() { close(toolEntered) })
			<-release
			return &iop.WorkspaceToolResponse{RequestId: req.GetRequestId(), StageId: req.GetStageId(), ToolCallId: req.GetToolCallId(), Status: iop.WorkspaceStatus_WORKSPACE_STATUS_CANCELLED, ErrorCode: iop.WorkspaceErrorCode_WORKSPACE_ERROR_CODE_CANCELLED, Error: "workspace command cancelled"}
		}
		execution := startWorkCoordinator(t, context.Background(), harness)
		select {
		case <-toolEntered:
		case <-time.After(2 * time.Second):
			close(release)
			t.Fatal("deadline tool did not reach Node")
		}
		_, err := waitWorkExecution(t, execution)
		close(release)
		outcome := waitWorkOutcome(t, harness.executor)
		if !errors.Is(err, edgeservice.ErrSingleRequestInternalToolBudget) || !errors.Is(outcome.err, errSingleRequestWorkStage) {
			t.Fatalf("wait err=%v outcome=%v", err, outcome.err)
		}
		if providerCalls.Load() != 1 || harness.node.toolCount.Load() != 1 || harness.executor.continueCount.Load() != 0 || harness.node.cancelCount.Load() > 1 || harness.node.cleanupCount.Load() != 1 || harness.bridge.pendingCount() != 0 {
			t.Fatalf("provider=%d tool=%d continuations=%d cancel=%d pending=%d", providerCalls.Load(), harness.node.toolCount.Load(), harness.executor.continueCount.Load(), harness.node.cancelCount.Load(), harness.bridge.pendingCount())
		}
	})
}

func TestSingleRequestWorkStageCancellation(t *testing.T) {
	t.Run("provider wait", func(t *testing.T) {
		frames := make(chan *iop.ProviderTunnelFrame)
		providerEntered := make(chan struct{})
		var providerCalls atomic.Int32
		provider := &mockService{submit: func(context.Context, edgeservice.ProviderPoolDispatchRequest) (*edgeservice.ProviderPoolDispatchResult, error) {
			providerCalls.Add(1)
			close(providerEntered)
			return workProviderDispatchResult(frames), nil
		}}
		ctx, cancel := context.WithCancel(context.Background())
		resultCh := make(chan struct {
			bridge *singleRequestWorkToolBridge
			err    error
		}, 1)
		go func() {
			bridge, err := runStandaloneWorkStageForTest(t, ctx, provider)
			resultCh <- struct {
				bridge *singleRequestWorkToolBridge
				err    error
			}{bridge: bridge, err: err}
		}()
		select {
		case <-providerEntered:
		case <-time.After(2 * time.Second):
			t.Fatal("provider wait did not start")
		}
		cancel()
		select {
		case result := <-resultCh:
			if !errors.Is(result.err, errSingleRequestWorkStage) || providerCalls.Load() != 1 || result.bridge.pendingCount() != 0 {
				t.Fatalf("err=%v provider=%d pending=%d", result.err, providerCalls.Load(), result.bridge.pendingCount())
			}
		case <-time.After(2 * time.Second):
			t.Fatal("provider cancellation did not return")
		}
	})

	t.Run("tool wait", func(t *testing.T) {
		var providerCalls atomic.Int32
		provider := scriptedWorkProvider([][]byte{
			workToolBody("cancel-1", edgeservice.InternalWorkspaceToolCommand, `{"command_id":"verify"}`),
		}, &providerCalls)
		harness := newWorkCoordinatorHarness(t, provider, nil)
		toolEntered := make(chan struct{})
		release := make(chan struct{})
		var entered sync.Once
		harness.node.toolResponder = func(req *iop.WorkspaceToolRequest) *iop.WorkspaceToolResponse {
			entered.Do(func() { close(toolEntered) })
			<-release
			return &iop.WorkspaceToolResponse{RequestId: req.GetRequestId(), StageId: req.GetStageId(), ToolCallId: req.GetToolCallId(), Status: iop.WorkspaceStatus_WORKSPACE_STATUS_CANCELLED, ErrorCode: iop.WorkspaceErrorCode_WORKSPACE_ERROR_CODE_CANCELLED, Error: "workspace command cancelled"}
		}
		ctx, cancel := context.WithCancel(context.Background())
		execution := startWorkCoordinator(t, ctx, harness)
		select {
		case <-toolEntered:
		case <-time.After(2 * time.Second):
			close(release)
			t.Fatal("tool wait did not reach Node")
		}
		cancel()
		select {
		case cancelRequest := <-harness.node.cancelRequests:
			if cancelRequest.GetRequestId() != "request-work" || cancelRequest.GetStageId() != singleRequestWorkStageID || cancelRequest.GetToolCallId() != "cancel-1" {
				close(release)
				t.Fatalf("cancel request=%+v", cancelRequest)
			}
		case <-time.After(2 * time.Second):
			close(release)
			t.Fatal("typed tool cancel did not reach Node")
		}
		close(release)
		_, err := waitWorkExecution(t, execution)
		outcome := waitWorkOutcome(t, harness.executor)
		if !errors.Is(err, edgeservice.ErrSingleRequestCancelled) || !errors.Is(outcome.err, errSingleRequestWorkStage) {
			t.Fatalf("wait err=%v outcome=%v", err, outcome.err)
		}
		time.Sleep(20 * time.Millisecond)
		if providerCalls.Load() != 1 || harness.node.toolCount.Load() != 1 || harness.executor.continueCount.Load() != 0 || harness.node.cancelCount.Load() != 1 || harness.node.cleanupCount.Load() != 1 || harness.bridge.pendingCount() != 0 {
			t.Fatalf("provider=%d tool=%d continuations=%d cancel=%d cleanup=%d pending=%d", providerCalls.Load(), harness.node.toolCount.Load(), harness.executor.continueCount.Load(), harness.node.cancelCount.Load(), harness.node.cleanupCount.Load(), harness.bridge.pendingCount())
		}
	})
}

func runStandaloneWorkStageForTest(t *testing.T, ctx context.Context, provider edgeserviceRunner) (*singleRequestWorkToolBridge, error) {
	t.Helper()
	bridge := newSingleRequestWorkToolBridge()
	controller := &workController{binding: workBinding(t), plan: []byte("# Plan\n\n## Goal\nUpdate.\n\n## Steps\n- [P1] Write.\n- [P2] Verify.\n\n## Verification\n- Verify.\n"), bridge: bridge}
	err := newSingleRequestWorkStage(newSingleRequestProviderStage(provider), bridge).run(ctx, workRequest(), controller)
	return bridge, err
}

func TestSingleRequestWorkStageRejectsMismatchedStoredPlan(t *testing.T) {
	t.Run("plan with altered heading fails before provider dispatch", func(t *testing.T) {
		var calls atomic.Int32
		provider := &mockService{submit: func(context.Context, edgeservice.ProviderPoolDispatchRequest) (*edgeservice.ProviderPoolDispatchResult, error) {
			calls.Add(1)
			return nil, errors.New("should not dispatch")
		}}
		bridge := newSingleRequestWorkToolBridge()
		ctrl := &workController{
			binding: workBinding(t),
			plan:    []byte("# Plan Mismatch\n\n## Goal\nUpdate.\n\n## Steps\n- [P1] Write.\n- [P2] Verify.\n\n## Verification\n- Verify.\n"),
			bridge: bridge,
		}
		err := newSingleRequestWorkStage(newSingleRequestProviderStage(provider), bridge).run(context.Background(), workRequest(), ctrl)
		if !errors.Is(err, errSingleRequestWorkStage) {
			t.Fatalf("err=%v, want errSingleRequestWorkStage", err)
		}
		if calls.Load() != 0 {
			t.Fatalf("provider dispatches=%d, want 0", calls.Load())
		}
		if len(ctrl.review) != 0 {
			t.Fatalf("review artifact written=%q, want empty", ctrl.review)
		}
		if bridge.pendingCount() != 0 {
			t.Fatalf("pending=%d, want 0", bridge.pendingCount())
		}
	})

	t.Run("plan with unknown placeholder fails before provider dispatch", func(t *testing.T) {
		var calls atomic.Int32
		provider := &mockService{submit: func(context.Context, edgeservice.ProviderPoolDispatchRequest) (*edgeservice.ProviderPoolDispatchResult, error) {
			calls.Add(1)
			return nil, errors.New("should not dispatch")
		}}
		bridge := newSingleRequestWorkToolBridge()
		ctrl := &workController{
			binding: workBinding(t),
			plan:    []byte("# Plan\n\n## Goal\nUpdate.\n\n## Steps\n- [P1] Write.\n- [P2] Verify.\n\n## Verification\n- Verify.\n{{unknown}}\n"),
			bridge: bridge,
		}
		err := newSingleRequestWorkStage(newSingleRequestProviderStage(provider), bridge).run(context.Background(), workRequest(), ctrl)
		if !errors.Is(err, errSingleRequestWorkStage) {
			t.Fatalf("err=%v, want errSingleRequestWorkStage", err)
		}
		if calls.Load() != 0 {
			t.Fatalf("provider dispatches=%d, want 0", calls.Load())
		}
	})
}

func TestSingleRequestWorkStageDrivesOrderedToolLoop(t *testing.T) {
	responses := [][]byte{
		workToolBody("write-1", edgeservice.InternalWorkspaceToolWrite, `{"relative_path":"result.txt","content":"done"}`),
		workToolBody("verify-1", edgeservice.InternalWorkspaceToolCommand, `{"command_id":"verify"}`),
		successBody(`{"item_status":"- P1: completed\n- P2: completed","changes":"Changed result.txt.","verification":"verify passed","deviations":"None"}`),
	}
	var bodies [][]byte
	var mu sync.Mutex
	provider := newSingleRequestProviderStage(&mockService{submit: func(_ context.Context, req edgeservice.ProviderPoolDispatchRequest) (*edgeservice.ProviderPoolDispatchResult, error) {
		body, err := req.Tunnel.BuildBody("ornith-fast")
		if err != nil {
			return nil, err
		}
		mu.Lock()
		bodies = append(bodies, body)
		index := len(bodies) - 1
		mu.Unlock()
		if index >= len(responses) {
			return nil, errors.New("unexpected provider call")
		}
		return &edgeservice.ProviderPoolDispatchResult{Path: edgeservice.ProviderPoolPathTunnel, Tunnel: &mockTunnel{frames: framesFor(responses[index])}, DispatchInfo: edgeservice.RunDispatch{ModelGroupKey: "ornith-fast", ProviderID: "gemini", Target: "ornith-fast", ProfileID: "profile-1", ProfileDriver: string(config.ProtocolDriverOpenAIChat), CredentialSlotRef: "slot-1", CredentialRevision: 1, ExecutionPath: string(edgeservice.ProviderPoolPathTunnel)}}, nil
	}})
	bridge := newSingleRequestWorkToolBridge()
	ctrl := &workController{binding: workBinding(t), plan: []byte("# Plan\n\n## Goal\nUpdate.\n\n## Steps\n- [P1] Write.\n- [P2] Verify.\n\n## Verification\n- Verify.\n"), bridge: bridge}
	err := newSingleRequestWorkStage(provider, bridge).run(context.Background(), workRequest(), ctrl)
	if err != nil {
		t.Fatal(err)
	}
	if !strings.Contains(string(ctrl.review), "Changed result.txt.") || !strings.Contains(string(ctrl.review), "verify passed") {
		t.Fatalf("handoff=%q", ctrl.review)
	}
	if bridge.pendingCount() != 0 || len(ctrl.envelopes) != 5 {
		t.Fatalf("pending=%d envelopes=%+v", bridge.pendingCount(), ctrl.envelopes)
	}
	for i, body := range bodies {
		if containsAll(string(body), "reasoning_effort") {
			t.Fatalf("body %d inherited reasoning: %s", i, body)
		}
		if !containsAll(string(body), "ornith-fast", "workspace_write", "workspace_command") {
			t.Fatalf("body %d missing admitted tools: %s", i, body)
		}
		var decoded map[string]any
		if err := json.Unmarshal(body, &decoded); err != nil {
			t.Fatalf("decode body %d: %v", i, err)
		}
		wantChoice := "auto"
		if i == 0 {
			wantChoice = "required"
		}
		if decoded["tool_choice"] != wantChoice {
			t.Fatalf("body %d tool_choice=%v, want %s", i, decoded["tool_choice"], wantChoice)
		}
		_, hasResponseFormat := decoded["response_format"]
		if hasResponseFormat != (i > 0) {
			t.Fatalf("body %d response_format present=%v, want %v", i, hasResponseFormat, i > 0)
		}
	}
	if !containsAll(string(bodies[0]), "PLAN", "[P1] Write.") || !containsAll(string(bodies[1]), "write-1", "verified") || !containsAll(string(bodies[2]), "verify-1", "verified") {
		t.Fatalf("tool continuation messages missing: %q", bodies)
	}
}

func TestSingleRequestWorkStageAdmitsCanonicalIOPDiagnostics(t *testing.T) {
	raw := []byte(`{"choices":[{"finish_reason":"tool_calls","index":0,"message":{"role":"assistant","content":"","reasoning_content":"private reasoning","tool_calls":[{"id":"write-1","type":"function","function":{"name":"workspace_write","arguments":"{\"relative_path\":\"smoke-result.txt\",\"content\":\"verified\"}"}}]}}],"created":1,"model":"ornith-fast","system_fingerprint":"b9193-test","object":"chat.completion","usage":{"completion_tokens":12,"prompt_tokens":34,"total_tokens":46},"id":"chatcmpl-test","timings":{"prompt_n":34,"prompt_ms":10.5,"predicted_n":12,"predicted_ms":20.5}}`)

	got, err := decodeSingleRequestWorkProviderResponse(raw, 4096)
	if err != nil {
		t.Fatal(err)
	}
	if got == nil || got.completion != nil || got.call == nil {
		t.Fatalf("response=%+v", got)
	}
	if got.call.ID != "write-1" || got.call.Type != "function" || got.call.Function.Name != edgeservice.InternalWorkspaceToolWrite || got.call.Function.Arguments != `{"relative_path":"smoke-result.txt","content":"verified"}` {
		t.Fatalf("call=%+v", got.call)
	}
}

func TestSingleRequestWorkStageRejectsMalformedResponsesAndOptions(t *testing.T) {
	bad := []string{
		`{"choices":[]}`,
		`{"id":"x","object":"chat.completion","created":1,"model":"m","unexpected":true,"choices":[]}`,
		`{"id":"x","object":"chat.completion","created":1,"model":"m","choices":[{"index":0,"finish_reason":"tool_calls","message":{"role":"assistant","content":"visible content is incompatible with a tool call","tool_calls":[{"id":"a","type":"function","function":{"name":"workspace_read","arguments":"{\"relative_path\":\"a\"}"}}]}}]}`,
		`{"id":"x","object":"chat.completion","created":1,"model":"m","choices":[{"index":0,"finish_reason":"stop","message":{"role":"assistant","content":"{\"completion\":\"c\",\"verification\":\"v\"}","unexpected":true}}]}`,
		`{"id":"x","object":"chat.completion","created":1,"model":"m","choices":[{"index":0,"finish_reason":"stop","message":{"role":"assistant","content":"{\"completion\":\"c\",\"verification\":\"v\"}","reasoning_content":7}}]}`,
		`{"id":"x","object":"chat.completion","created":1,"model":"m","choices":[{"index":0,"finish_reason":"stop","message":{"role":"assistant","content":"{\"completion\":\"c\",\"verification\":\"v\"}","extra_content":{"google":{"thought_signature":"not-admitted-for-work"}}}}]}`,
		`{"id":"x","object":"chat.completion","created":1,"model":"m","choices":[{"index":0,"finish_reason":"tool_calls","message":{"role":"assistant","content":null,"tool_calls":[]}}]}`,
		`{"id":"x","object":"chat.completion","created":1,"model":"m","choices":[{"index":0,"finish_reason":"tool_calls","message":{"role":"assistant","content":null,"tool_calls":[{"id":"a","type":"function","function":{"name":"workspace_read","arguments":"{}"}},{"id":"b","type":"function","function":{"name":"workspace_read","arguments":"{}"}}]}}]}`,
		`{"id":"x","object":"chat.completion","created":1,"model":"m","choices":[{"index":0,"finish_reason":"stop","message":{"role":"assistant","content":"{\"completion\":\"\",\"verification\":\"v\"}"}}]}`,
		`{"id":"x","object":"chat.completion","created":1,"model":"m","choices":[{"index":0,"finish_reason":"stop","message":{"role":"assistant","content":"{\"Completion\":\"c\",\"verification\":\"v\"}"}}]}`,
	}
	for _, raw := range bad {
		_, err := decodeSingleRequestWorkProviderResponse([]byte(raw), 4096)
		if !errors.Is(err, errSingleRequestWorkStage) {
			t.Fatalf("raw=%s err=%v", raw, err)
		}
	}
	request := workRequest()
	request.StageBinding.Options["reasoning_effort"] = "high"
	bridge := newSingleRequestWorkToolBridge()
	ctrl := &workController{binding: workBinding(t), plan: []byte("# Plan\n\n## Goal\nUpdate.\n\n## Steps\n- [P1] Write.\n- [P2] Verify.\n\n## Verification\n- Verify.\n"), bridge: bridge}
	if err := newSingleRequestWorkStage(newSingleRequestProviderStage(&mockService{}), bridge).run(context.Background(), request, ctrl); !errors.Is(err, errSingleRequestWorkStage) {
		t.Fatalf("err=%v", err)
	}
	for name, arguments := range map[string]string{
		"non-object":         `[]`,
		"duplicate field":    `{"relative_path":"a","relative_path":"b"}`,
		"trailing value":     `{"relative_path":"a"} {}`,
		"quoted nested json": `"{\"relative_path\":\"a\"}"`,
	} {
		t.Run(name, func(t *testing.T) {
			if _, err := decodeSingleRequestWorkToolArguments(arguments); !errors.Is(err, errSingleRequestWorkStage) {
				t.Fatalf("decode error=%v", err)
			}
		})
	}
}

func TestNormalizeSingleRequestProviderToolArguments(t *testing.T) {
	tests := []struct {
		name      string
		tool      string
		arguments string
		want      string
	}{
		{name: "empty list root becomes canonical dot", tool: edgeservice.InternalWorkspaceToolList, arguments: `{"relative_path":""}`, want: `{"relative_path":"."}`},
		{name: "empty read remains invalid for service", tool: edgeservice.InternalWorkspaceToolRead, arguments: `{"relative_path":""}`, want: `{"relative_path":""}`},
		{name: "dot-prefixed list remains invalid for service", tool: edgeservice.InternalWorkspaceToolList, arguments: `{"relative_path":"./"}`, want: `{"relative_path":"./"}`},
		{name: "canonical list path is unchanged", tool: edgeservice.InternalWorkspaceToolList, arguments: `{"relative_path":"src"}`, want: `{"relative_path":"src"}`},
		{name: "unknown list field is unchanged", tool: edgeservice.InternalWorkspaceToolList, arguments: `{"relative_path":"","extra":true}`, want: `{"relative_path":"","extra":true}`},
	}
	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			got := normalizeSingleRequestProviderToolArguments(test.tool, json.RawMessage(test.arguments))
			if string(got) != test.want {
				t.Fatalf("normalized arguments=%s, want %s", got, test.want)
			}
		})
	}
}

func TestSingleRequestWorkStageRejectsReservedOptionAliases(t *testing.T) {
	messageSets := map[string][]chatMessage{
		"initial": {{Role: "user", Content: "immutable task"}},
		"resumed": {
			{Role: "assistant", ToolCalls: []any{map[string]any{"id": "tool-1"}}},
			{Role: "tool", ToolCallID: "tool-1", ToolName: edgeservice.InternalWorkspaceToolRead, Content: `{}`},
		},
	}
	aliases := map[string]string{
		"model":               "Model",
		"messages":            "MESSAGES",
		"tools":               "Tools",
		"tool_choice":         "Tool_Choice",
		"parallel_tool_calls": "Parallel_Tool_Calls",
		"response_format":     "Response_Format",
		"stream":              "Stream",
		"credential":          "Credential",
		"credential_binding":  "Credential_Binding",
		"reasoning_effort":    "Reasoning_Effort",
	}
	tools := []any{singleRequestWorkToolSchema(edgeservice.InternalWorkspaceToolRead, map[string]any{"type": "object"})}
	for messageName, messages := range messageSets {
		completionEligible := messageName == "resumed"
		for canonical, alias := range aliases {
			t.Run(messageName+"/"+canonical, func(t *testing.T) {
				if _, err := buildSingleRequestWorkBody(messages, map[string]any{alias: "forbidden"}, tools, "ornith-fast", completionEligible); !errors.Is(err, errSingleRequestWorkStage) {
					t.Fatalf("alias %q error=%v", alias, err)
				}
			})
		}
	}

	body, err := buildSingleRequestWorkBody(messageSets["initial"], map[string]any{
		"model": "override", "messages": "override", "tools": "override", "tool_choice": "required",
		"parallel_tool_calls": true, "response_format": "override", "stream": true, "credential": "secret", "credential_binding": "private",
		"temperature": 0.1,
	}, tools, "ornith-fast", false)
	if err != nil {
		t.Fatal(err)
	}
	var decoded map[string]any
	if json.Unmarshal(body, &decoded) != nil || decoded["model"] != "ornith-fast" || decoded["tool_choice"] != "required" || decoded["parallel_tool_calls"] != false || decoded["stream"] != false || decoded["temperature"] != 0.1 {
		t.Fatalf("server-owned body=%s", body)
	}
	if _, exists := decoded["credential"]; exists {
		t.Fatalf("credential serialized: %s", body)
	}
	if _, exists := decoded["credential_binding"]; exists {
		t.Fatalf("credential binding serialized: %s", body)
	}
	if _, exists := decoded["response_format"]; exists {
		t.Fatalf("completion format serialized while tool-required: %s", body)
	}
	resumedBody, err := buildSingleRequestWorkBody(messageSets["resumed"], nil, tools, "ornith-fast", true)
	if err != nil {
		t.Fatal(err)
	}
	if json.Unmarshal(resumedBody, &decoded) != nil || decoded["tool_choice"] != "auto" {
		t.Fatalf("resumed server-owned body=%s", resumedBody)
	}
	encodedFormat, err := json.Marshal(singleRequestWorkResponseFormat())
	if err != nil {
		t.Fatal(err)
	}
	var expectedFormat any
	if json.Unmarshal(encodedFormat, &expectedFormat) != nil || !jsonValuesEqual(decoded["response_format"], expectedFormat) {
		t.Fatalf("response format not server-owned: %s", resumedBody)
	}
}

func TestSingleRequestWorkStageProjectsClosedEnvironmentSchema(t *testing.T) {
	workspace := workBinding(t).Workspace.Clone()
	workspace.EnvironmentNames = []string{"SAFE"}
	tools, err := singleRequestWorkTools(workspace)
	if err != nil {
		t.Fatal(err)
	}
	var environment map[string]any
	for _, raw := range tools {
		tool, ok := raw.(map[string]any)
		if !ok {
			continue
		}
		function, _ := tool["function"].(map[string]any)
		if function["name"] != edgeservice.InternalWorkspaceToolCommand {
			continue
		}
		parameters, _ := function["parameters"].(map[string]any)
		properties, _ := parameters["properties"].(map[string]any)
		environment, _ = properties["environment"].(map[string]any)
	}
	if environment == nil || environment["additionalProperties"] != false {
		t.Fatalf("environment schema=%+v", environment)
	}
	properties, ok := environment["properties"].(map[string]any)
	if !ok || len(properties) != 1 || properties["SAFE"] == nil || properties["NOT_ALLOWED"] != nil {
		t.Fatalf("environment properties=%+v", environment["properties"])
	}
}

func TestSingleRequestWorkspaceToolSchemasDescribeCanonicalRelativePaths(t *testing.T) {
	tools, err := singleRequestWorkTools(workBinding(t).Workspace)
	if err != nil {
		t.Fatal(err)
	}
	pathTools := 0
	for _, raw := range tools {
		tool, _ := raw.(map[string]any)
		function, _ := tool["function"].(map[string]any)
		if function["strict"] != true {
			t.Fatalf("tool %q strict=%v, want true", function["name"], function["strict"])
		}
		parameters, _ := function["parameters"].(map[string]any)
		properties, _ := parameters["properties"].(map[string]any)
		pathProperty, hasPath := properties["relative_path"].(map[string]any)
		if !hasPath {
			continue
		}
		pathTools++
		if pathProperty["description"] != singleRequestCanonicalRelativePathDescription {
			t.Fatalf("tool %q relative_path description=%v", function["name"], pathProperty["description"])
		}
		if pathProperty["minLength"] != 1 {
			t.Fatalf("tool %q relative_path minLength=%v, want 1", function["name"], pathProperty["minLength"])
		}
	}
	if pathTools == 0 {
		t.Fatal("workspace tools did not expose a relative_path schema")
	}
}

func TestSingleRequestWorkToolBridgeCorrelatesAndCleansUp(t *testing.T) {
	b := newSingleRequestWorkToolBridge()
	keys := []singleRequestWorkToolKey{{"request-a", "working", "one"}, {"request-b", "working", "two"}}
	channels := make([]<-chan edgeservice.InternalWorkspaceToolResult, len(keys))
	for i, key := range keys {
		var err error
		channels[i], err = b.register(key)
		if err != nil {
			t.Fatal(err)
		}
	}
	if _, err := b.register(keys[0]); !errors.Is(err, errSingleRequestWorkStage) {
		t.Fatalf("duplicate err=%v", err)
	}
	if err := b.ContinueInternalTool(context.Background(), edgeservice.InternalWorkspaceToolResult{RequestID: "request-b", StageID: "working", ToolCallID: "two", Status: "success"}); err != nil {
		t.Fatal(err)
	}
	if err := b.ContinueInternalTool(context.Background(), edgeservice.InternalWorkspaceToolResult{RequestID: "request-a", StageID: "working", ToolCallID: "one", Status: "success"}); err != nil {
		t.Fatal(err)
	}
	for i, key := range keys {
		result, err := b.wait(context.Background(), key, channels[i])
		if err != nil || result.ToolCallID != key.toolCallID {
			t.Fatalf("key=%+v result=%+v err=%v", key, result, err)
		}
	}
	if b.pendingCount() != 0 {
		t.Fatalf("pending=%d", b.pendingCount())
	}
	ch, err := b.register(singleRequestWorkToolKey{"request-c", "working", "cancel"})
	if err != nil {
		t.Fatal(err)
	}
	ctx, cancel := context.WithCancel(context.Background())
	cancel()
	if _, err := b.wait(ctx, singleRequestWorkToolKey{"request-c", "working", "cancel"}, ch); !errors.Is(err, errSingleRequestWorkStage) {
		t.Fatalf("cancel err=%v", err)
	}
	if b.pendingCount() != 0 {
		t.Fatalf("pending after cancel=%d", b.pendingCount())
	}
	if err := b.ContinueInternalTool(context.Background(), edgeservice.InternalWorkspaceToolResult{RequestID: "request-c", StageID: "working", ToolCallID: "cancel"}); !errors.Is(err, errSingleRequestWorkStage) {
		t.Fatalf("stale err=%v", err)
	}
}

func TestSingleRequestWorkToolBridgeRace(t *testing.T) {
	b := newSingleRequestWorkToolBridge()
	const workers = 32
	var group sync.WaitGroup
	for i := 0; i < workers; i++ {
		group.Add(1)
		go func(i int) {
			defer group.Done()
			key := singleRequestWorkToolKey{requestID: "request", stageID: "working", toolCallID: string(rune('a' + i))}
			ch, err := b.register(key)
			if err != nil {
				t.Errorf("register: %v", err)
				return
			}
			if err := b.ContinueInternalTool(context.Background(), edgeservice.InternalWorkspaceToolResult{RequestID: key.requestID, StageID: key.stageID, ToolCallID: key.toolCallID}); err != nil {
				t.Errorf("continue: %v", err)
				return
			}
			if _, err := b.wait(context.Background(), key, ch); err != nil {
				t.Errorf("wait: %v", err)
			}
		}(i)
	}
	group.Wait()
	if b.pendingCount() != 0 {
		t.Fatalf("pending=%d", b.pendingCount())
	}
}

func TestSingleRequestWorkToolBridgeRejectsConcurrentDuplicate(t *testing.T) {
	b := newSingleRequestWorkToolBridge()
	key := singleRequestWorkToolKey{requestID: "request", stageID: singleRequestWorkStageID, toolCallID: "duplicate"}
	resultCh, err := b.register(key)
	if err != nil {
		t.Fatal(err)
	}
	start := make(chan struct{})
	errorsCh := make(chan error, 2)
	var group sync.WaitGroup
	for range 2 {
		group.Add(1)
		go func() {
			defer group.Done()
			<-start
			errorsCh <- b.ContinueInternalTool(context.Background(), edgeservice.InternalWorkspaceToolResult{RequestID: key.requestID, StageID: key.stageID, ToolCallID: key.toolCallID, Status: "success"})
		}()
	}
	close(start)
	group.Wait()
	close(errorsCh)
	successes, rejections := 0, 0
	for err := range errorsCh {
		if err == nil {
			successes++
		} else if errors.Is(err, errSingleRequestWorkStage) {
			rejections++
		} else {
			t.Fatalf("unexpected delivery error: %v", err)
		}
	}
	if successes != 1 || rejections != 1 || b.pendingCount() != 0 {
		t.Fatalf("successes=%d rejections=%d pending=%d", successes, rejections, b.pendingCount())
	}
	result, err := b.wait(context.Background(), key, resultCh)
	if err != nil || result.ToolCallID != key.toolCallID || b.pendingCount() != 0 {
		t.Fatalf("result=%+v err=%v pending=%d", result, err, b.pendingCount())
	}
}
