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

package service

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

	toki "git.toki-labs.com/toki/proto-socket/go"

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

type scriptedInternalToolExecutor struct {
	calls         []InternalWorkspaceToolCall
	results       chan InternalWorkspaceToolResult
	seenResults   []InternalWorkspaceToolResult
	continueCount atomic.Int32
}

func newScriptedInternalToolExecutor(calls ...InternalWorkspaceToolCall) *scriptedInternalToolExecutor {
	return &scriptedInternalToolExecutor{calls: calls, results: make(chan InternalWorkspaceToolResult, len(calls)+1)}
}

func (e *scriptedInternalToolExecutor) ExecuteSingleRequest(ctx context.Context, req SingleRequestRequest, ctrl SingleRequestController) error {
	sequence := uint64(1)
	if err := ctrl.SubmitEnvelope(testEnvelope(req.RequestID, sequence, SingleRequestStatePlanning)); err != nil {
		return err
	}
	for index := range e.calls {
		call := e.calls[index].Clone()
		if call.RequestID == "" {
			call.RequestID = req.RequestID
		}
		if call.StageID == "" {
			call.StageID = "plan"
		}
		sequence++
		if err := ctrl.SubmitEnvelope(SingleRequestEnvelope{
			RequestID: req.RequestID, Sequence: sequence,
			Stage: SingleRequestStateInternalTool, SavedStage: SingleRequestStatePlanning,
			ToolCall: call,
		}); err != nil {
			return err
		}
		select {
		case result := <-e.results:
			e.seenResults = append(e.seenResults, result.Clone())
		case <-ctx.Done():
			return ctx.Err()
		}
		sequence++
		if err := ctrl.SubmitEnvelope(SingleRequestEnvelope{
			RequestID: req.RequestID, Sequence: sequence,
			Stage: SingleRequestStatePlanning, SavedStage: SingleRequestStatePlanning,
		}); err != nil {
			return err
		}
	}
	for _, stage := range []SingleRequestState{SingleRequestStateWorking, SingleRequestStateReviewing, SingleRequestStateFinalizing} {
		sequence++
		envelope := testEnvelope(req.RequestID, sequence, stage)
		if stage == SingleRequestStateFinalizing {
			envelope.Result = &SingleRequestResult{Output: "private tools completed"}
		}
		if err := ctrl.SubmitEnvelope(envelope); err != nil {
			return err
		}
	}
	return nil
}

func (e *scriptedInternalToolExecutor) ContinueInternalTool(_ context.Context, result InternalWorkspaceToolResult) error {
	e.continueCount.Add(1)
	e.results <- result.Clone()
	return nil
}

func internalLoopWorkspace() config.WorkspaceDefinition {
	return config.WorkspaceDefinition{
		Ref: "workspace-loop", Platform: "darwin", Root: "/Users/operator/project",
		Operations: []config.WorkspaceOperation{
			config.WorkspaceOpRead, config.WorkspaceOpList, config.WorkspaceOpWrite,
			config.WorkspaceOpDelete, config.WorkspaceOpCommand,
		},
		Commands:             []config.WorkspaceCommandDefinition{{ID: "test"}},
		EnvironmentAllowlist: []string{"IOP_MODE"},
		MaxReadBytes:         1024, MaxWriteBytes: 1024, MaxOutputBytes: 1024, MaxCommandTimeoutMS: 1000,
	}
}

func newInternalToolLoopService(t *testing.T, executor SingleRequestExecutor) (*Service, *toki.TcpClient) {
	t.Helper()
	edgeClient, nodeClient := workspaceWirePipe(t)
	registry := edgenode.NewRegistry()
	registry.Register(&edgenode.NodeEntry{NodeID: "node-loop", Client: edgeClient})
	service := New(registry, nil)
	service.SetNodeStore(workspaceStore("node-loop", internalLoopWorkspace()))
	service.SetSingleRequestExecutor(executor)
	return service, nodeClient
}

func TestSingleRequestInternalToolLoopRequiresOptionalContinuation(t *testing.T) {
	executor := &channelFakeExecutor{fn: func(_ context.Context, req SingleRequestRequest, ctrl SingleRequestController) error {
		if err := ctrl.SubmitEnvelope(testEnvelope(req.RequestID, 1, SingleRequestStatePlanning)); err != nil {
			return err
		}
		return ctrl.SubmitEnvelope(SingleRequestEnvelope{
			RequestID: req.RequestID, Sequence: 2,
			Stage: SingleRequestStateInternalTool, SavedStage: SingleRequestStatePlanning,
			ToolCall: &InternalWorkspaceToolCall{
				RequestID: req.RequestID, StageID: "plan", ToolCallID: "tool-1",
				Name: InternalWorkspaceToolRead, Arguments: json.RawMessage(`{"relative_path":"README.md"}`),
			},
		})
	}}
	service, node := newInternalToolLoopService(t, executor)
	var openCount atomic.Int32
	installInternalLoopOpenResponder(node, &openCount)
	handle, err := service.StartSingleRequest(context.Background(), internalLoopRequest(t, nil))
	if err != nil {
		t.Fatalf("StartSingleRequest: %v", err)
	}
	if _, err := waitForExecution(t, handle); !errors.Is(err, ErrSingleRequestInternalToolUnavailable) {
		t.Fatalf("Wait error = %v, want unavailable continuation", err)
	}
	if openCount.Load() != 0 {
		t.Fatalf("unavailable continuation opened workspace %d times", openCount.Load())
	}
}

func internalLoopRequest(t *testing.T, mutate func(*SingleRequestBinding)) SingleRequestRequest {
	t.Helper()
	binding := createTestBinding(t)
	binding.WorkspaceRef = "workspace-loop"
	if mutate != nil {
		mutate(binding)
	}
	return SingleRequestRequest{RequestID: "request-loop", Binding: binding, Prompt: "complete the private task"}
}

func installInternalLoopOpenResponder(node *toki.TcpClient, count *atomic.Int32) {
	toki.AddRequestListenerTyped[*iop.WorkspaceOpenRequest, *iop.WorkspaceOpenResponse](&node.Communicator, func(req *iop.WorkspaceOpenRequest) (*iop.WorkspaceOpenResponse, error) {
		count.Add(1)
		return &iop.WorkspaceOpenResponse{
			RequestId: req.GetRequestId(), WorkspaceRef: req.GetWorkspaceRef(),
			Status: iop.WorkspaceStatus_WORKSPACE_STATUS_SUCCESS,
		}, nil
	})
	toki.AddRequestListenerTyped[*iop.WorkspaceCleanupRequest, *iop.WorkspaceCleanupResponse](&node.Communicator, func(req *iop.WorkspaceCleanupRequest) (*iop.WorkspaceCleanupResponse, error) {
		return &iop.WorkspaceCleanupResponse{RequestId: req.GetRequestId(), Status: iop.WorkspaceStatus_WORKSPACE_STATUS_SUCCESS}, nil
	})
}

func TestSingleRequestInternalToolLoopMultipleOrdered(t *testing.T) {
	executor := newScriptedInternalToolExecutor(
		InternalWorkspaceToolCall{ToolCallID: "tool-read", Name: InternalWorkspaceToolRead, Arguments: json.RawMessage(`{"relative_path":"README.md"}`)},
		InternalWorkspaceToolCall{ToolCallID: "tool-write", Name: InternalWorkspaceToolWrite, Arguments: json.RawMessage(`{"relative_path":"result.txt","content":"done"}`)},
		InternalWorkspaceToolCall{ToolCallID: "tool-command", Name: InternalWorkspaceToolCommand, Arguments: json.RawMessage(`{"command_id":"test","environment":{"IOP_MODE":"safe"}}`)},
	)
	service, node := newInternalToolLoopService(t, executor)
	var openCount atomic.Int32
	installInternalLoopOpenResponder(node, &openCount)
	toolOrder := make(chan string, 3)
	commandTimeout := make(chan int64, 1)
	toki.AddRequestListenerTyped[*iop.WorkspaceToolRequest, *iop.WorkspaceToolResponse](&node.Communicator, func(req *iop.WorkspaceToolRequest) (*iop.WorkspaceToolResponse, error) {
		toolOrder <- req.GetToolCallId()
		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_READ:
			response.Content = []byte("source")
		case iop.WorkspaceOperation_WORKSPACE_OPERATION_COMMAND:
			commandTimeout <- req.GetTimeoutMs()
			response.Stdout = []byte("ok")
		}
		return response, nil
	})

	handle, err := service.StartSingleRequest(context.Background(), internalLoopRequest(t, nil))
	if err != nil {
		t.Fatalf("StartSingleRequest: %v", err)
	}
	waitForState(t, handle, SingleRequestStateFinalizing)
	waitForSingleRequestCleanup(t, handle)
	if err := handle.AcknowledgeTerminal(true); err != nil {
		t.Fatalf("AcknowledgeTerminal: %v", err)
	}
	result, err := waitForExecution(t, handle)
	if err != nil || result.Output != "private tools completed" {
		t.Fatalf("Wait = (%q, %v)", result.Output, err)
	}
	if openCount.Load() != 1 || executor.continueCount.Load() != 3 || len(executor.seenResults) != 3 {
		t.Fatalf("open=%d continuations=%d results=%d", openCount.Load(), executor.continueCount.Load(), len(executor.seenResults))
	}
	for index, want := range []string{"tool-read", "tool-write", "tool-command"} {
		if got := <-toolOrder; got != want {
			t.Fatalf("tool order[%d]=%q, want %q", index, got, want)
		}
		if executor.seenResults[index].ToolCallID != want || executor.seenResults[index].StageID != "plan" {
			t.Fatalf("correlated result[%d]=%+v", index, executor.seenResults[index])
		}
	}
	if got := <-commandTimeout; got <= 0 || got > int64(internalLoopWorkspace().MaxCommandTimeoutMS) {
		t.Fatalf("command timeout = %d, want within workspace maximum %d", got, internalLoopWorkspace().MaxCommandTimeoutMS)
	}
	for progress := range handle.Progress() {
		if progress.Message == InternalWorkspaceToolRead || progress.Message == InternalWorkspaceToolWrite || progress.Message == InternalWorkspaceToolCommand {
			t.Fatalf("internal tool protocol reached progress: %+v", progress)
		}
	}
}

func waitForSingleRequestCleanup(t *testing.T, handle SingleRequestExecution) {
	t.Helper()
	internal, ok := handle.(*singleRequestHandle)
	if !ok {
		t.Fatal("execution does not expose coordinator cleanup state")
	}
	deadline := time.Now().Add(2 * time.Second)
	for time.Now().Before(deadline) {
		internal.mu.Lock()
		complete := internal.cleanupComplete
		internal.mu.Unlock()
		if complete {
			return
		}
		time.Sleep(time.Millisecond)
	}
	t.Fatal("workspace cleanup did not complete")
}

func assertSingleRequestRequestBudgetOwnership(t *testing.T, handle SingleRequestExecution, observer *capturingObserver, wantErr error, forbidden ...string) {
	t.Helper()
	result, waitErr := waitForExecution(t, handle)
	if !errors.Is(waitErr, wantErr) {
		t.Fatalf("Wait error = %v, want %v", waitErr, wantErr)
	}
	if result.Output != "" || handle.State() != SingleRequestStateFailed {
		t.Fatalf("result/state = (%+v, %s), want empty failed result", result, handle.State())
	}

	wantTerminal := SingleRequestTerminalDisposition{Kind: SingleRequestTerminalError, ErrorClass: SingleRequestTerminalErrorBudget}
	terminalCount := 0
	for progress := range handle.Progress() {
		if progress.Terminal == nil {
			continue
		}
		terminalCount++
		if *progress.Terminal != wantTerminal {
			t.Fatalf("terminal = %+v, want %+v", *progress.Terminal, wantTerminal)
		}
	}
	if terminalCount != 1 {
		t.Fatalf("terminal count = %d, want exactly one", terminalCount)
	}

	events := observer.snapshot()
	assertSingleRequestCorrelation(t, events, forbidden...)
	terminalObservationCount := 0
	for _, event := range events {
		if event.ErrorClass == singleRequestErrorClassTimeout {
			t.Fatalf("request wall-clock expiry produced timeout observation: %#v", events)
		}
		if event.EventClass == singleRequestEventClassTerminal {
			terminalObservationCount++
			if event.Outcome != singleRequestOutcomeError || event.ErrorClass != singleRequestErrorClassInternalToolBudget {
				t.Fatalf("terminal observation = %#v, want error/internal_tool_budget", event)
			}
		}
	}
	if terminalObservationCount != 1 {
		t.Fatalf("terminal observation count = %d, want exactly one: %#v", terminalObservationCount, events)
	}
}

func TestSingleRequestInternalToolLoopFailsClosed(t *testing.T) {
	const rawSentinel = "RAW-TOOL-SENTINEL"
	tests := []struct {
		name              string
		calls             []InternalWorkspaceToolCall
		mutateBinding     func(*SingleRequestBinding)
		respond           func(*iop.WorkspaceToolRequest) *iop.WorkspaceToolResponse
		want              error
		wantWireCalls     int32
		wantContinuations int32
	}{
		{
			name:  "identity mismatch",
			calls: []InternalWorkspaceToolCall{{RequestID: "other-request", StageID: "plan", ToolCallID: "tool-1", Name: InternalWorkspaceToolRead, Arguments: json.RawMessage(`{"relative_path":"README.md"}`)}},
			want:  ErrSingleRequestIdentityMismatch,
		},
		{
			name:  "malformed arguments",
			calls: []InternalWorkspaceToolCall{{ToolCallID: "tool-1", Name: InternalWorkspaceToolRead, Arguments: json.RawMessage(`{"relative_path":"README.md","raw":"` + rawSentinel + `"}`)}},
			want:  ErrSingleRequestInternalToolInvalidCall,
		},
		{
			name:  "capability denied",
			calls: []InternalWorkspaceToolCall{{ToolCallID: "tool-1", Name: InternalWorkspaceToolCommand, Arguments: json.RawMessage(`{"command_id":"missing"}`)}},
			want:  ErrSingleRequestInternalToolDenied,
		},
		{
			name:  "environment capability denied",
			calls: []InternalWorkspaceToolCall{{ToolCallID: "tool-1", Name: InternalWorkspaceToolCommand, Arguments: json.RawMessage(`{"command_id":"test","environment":{"NOT_ALLOWED":"value"}}`)}},
			want:  ErrSingleRequestInternalToolDenied,
		},
		{
			name: "duplicate tool id",
			calls: []InternalWorkspaceToolCall{
				{ToolCallID: "tool-1", Name: InternalWorkspaceToolRead, Arguments: json.RawMessage(`{"relative_path":"README.md"}`)},
				{ToolCallID: "tool-1", Name: InternalWorkspaceToolRead, Arguments: json.RawMessage(`{"relative_path":"README.md"}`)},
			},
			want: ErrSingleRequestInternalToolInvalidCall, wantWireCalls: 1, wantContinuations: 1,
		},
		{
			name: "iteration budget",
			calls: []InternalWorkspaceToolCall{
				{ToolCallID: "tool-1", Name: InternalWorkspaceToolRead, Arguments: json.RawMessage(`{"relative_path":"README.md"}`)},
				{ToolCallID: "tool-2", Name: InternalWorkspaceToolRead, Arguments: json.RawMessage(`{"relative_path":"README.md"}`)},
			},
			mutateBinding: func(binding *SingleRequestBinding) { binding.Limits.MaxToolIterations = 1 },
			want:          ErrSingleRequestInternalToolBudget, wantWireCalls: 1, wantContinuations: 1,
		},
		{
			name:  "stale Node response",
			calls: []InternalWorkspaceToolCall{{ToolCallID: "tool-1", Name: InternalWorkspaceToolRead, Arguments: json.RawMessage(`{"relative_path":"README.md"}`)}},
			respond: func(req *iop.WorkspaceToolRequest) *iop.WorkspaceToolResponse {
				return &iop.WorkspaceToolResponse{RequestId: req.GetRequestId(), StageId: req.GetStageId(), ToolCallId: "stale-tool", Status: iop.WorkspaceStatus_WORKSPACE_STATUS_SUCCESS}
			},
			want: ErrSingleRequestInternalToolFailed, wantWireCalls: 1,
		},
		{
			name:          "output budget",
			calls:         []InternalWorkspaceToolCall{{ToolCallID: "tool-1", Name: InternalWorkspaceToolRead, Arguments: json.RawMessage(`{"relative_path":"README.md"}`)}},
			mutateBinding: func(binding *SingleRequestBinding) { binding.Limits.MaxOutputBytes = 3 },
			respond: 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("four")}
			},
			want: ErrSingleRequestInternalToolBudget, wantWireCalls: 1,
		},
		{
			name: "cumulative output budget",
			calls: []InternalWorkspaceToolCall{
				{ToolCallID: "tool-1", Name: InternalWorkspaceToolRead, Arguments: json.RawMessage(`{"relative_path":"README.md"}`)},
				{ToolCallID: "tool-2", Name: InternalWorkspaceToolRead, Arguments: json.RawMessage(`{"relative_path":"README.md"}`)},
			},
			mutateBinding: func(binding *SingleRequestBinding) { binding.Limits.MaxOutputBytes = 3 },
			respond: 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("xx")}
			},
			want: ErrSingleRequestInternalToolBudget, wantWireCalls: 2, wantContinuations: 1,
		},
	}
	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			executor := newScriptedInternalToolExecutor(test.calls...)
			service, node := newInternalToolLoopService(t, executor)
			var openCount atomic.Int32
			installInternalLoopOpenResponder(node, &openCount)
			var toolCount atomic.Int32
			toki.AddRequestListenerTyped[*iop.WorkspaceToolRequest, *iop.WorkspaceToolResponse](&node.Communicator, func(req *iop.WorkspaceToolRequest) (*iop.WorkspaceToolResponse, error) {
				toolCount.Add(1)
				if test.respond != nil {
					return test.respond(req), nil
				}
				return &iop.WorkspaceToolResponse{RequestId: req.GetRequestId(), StageId: req.GetStageId(), ToolCallId: req.GetToolCallId(), Status: iop.WorkspaceStatus_WORKSPACE_STATUS_SUCCESS}, nil
			})
			handle, err := service.StartSingleRequest(context.Background(), internalLoopRequest(t, test.mutateBinding))
			if err != nil {
				t.Fatalf("StartSingleRequest: %v", err)
			}
			_, err = waitForExecution(t, handle)
			if !errors.Is(err, test.want) {
				t.Fatalf("Wait error = %v, want %v", err, test.want)
			}
			if strings.Contains(err.Error(), rawSentinel) {
				t.Fatalf("raw call leaked in error %q", err)
			}
			wantOpenCalls := int32(0)
			if test.wantWireCalls > 0 {
				wantOpenCalls = 1
			}
			if openCount.Load() != wantOpenCalls || toolCount.Load() != test.wantWireCalls || executor.continueCount.Load() != test.wantContinuations {
				t.Fatalf("open=%d wire calls=%d continuations=%d, want %d/%d/%d", openCount.Load(), toolCount.Load(), executor.continueCount.Load(), wantOpenCalls, test.wantWireCalls, test.wantContinuations)
			}
		})
	}
}

func TestSingleRequestInternalToolLoopCancelPropagates(t *testing.T) {
	executor := newScriptedInternalToolExecutor(InternalWorkspaceToolCall{
		ToolCallID: "tool-command", Name: InternalWorkspaceToolCommand,
		Arguments: json.RawMessage(`{"command_id":"test"}`),
	})
	service, node := newInternalToolLoopService(t, executor)
	var openCount atomic.Int32
	installInternalLoopOpenResponder(node, &openCount)
	var sequence atomic.Int32
	toolEntered := make(chan struct{})
	release := make(chan struct{})
	cancelReached := make(chan *iop.WorkspaceCancelRequest, 1)
	serveWorkspaceConcurrent(&node.Communicator, &sequence, func(req *iop.WorkspaceToolRequest) *iop.WorkspaceToolResponse {
		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"}
	})
	serveWorkspaceConcurrent(&node.Communicator, &sequence, func(req *iop.WorkspaceCancelRequest) *iop.WorkspaceCancelResponse {
		cancelReached <- req
		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"}
	})

	handle, err := service.StartSingleRequest(context.Background(), internalLoopRequest(t, nil))
	if err != nil {
		t.Fatalf("StartSingleRequest: %v", err)
	}
	select {
	case <-toolEntered:
	case <-time.After(2 * time.Second):
		close(release)
		t.Fatal("tool did not reach Node")
	}
	handle.Cancel()
	select {
	case cancel := <-cancelReached:
		if cancel.GetRequestId() != "request-loop" || cancel.GetStageId() != "plan" || cancel.GetToolCallId() != "tool-command" {
			close(release)
			t.Fatalf("cancel identity = %+v", cancel)
		}
	case <-time.After(2 * time.Second):
		close(release)
		t.Fatal("typed cancel did not reach Node")
	}
	close(release)
	if _, err := waitForExecution(t, handle); !errors.Is(err, ErrSingleRequestCancelled) {
		t.Fatalf("Wait error = %v, want cancelled", err)
	}
	if executor.continueCount.Load() != 0 {
		t.Fatalf("cancelled call delivered %d continuations", executor.continueCount.Load())
	}
}

func TestSingleRequestLateInternalToolAdmissionCallerCancellation(t *testing.T) {
	now := time.Now()
	binding := createTestBinding(t)
	binding.Workspace = &SingleRequestWorkspaceBinding{
		Ref:          "workspace-loop",
		OperationIDs: []string{"read"},
		Limits:       SingleRequestWorkspaceLimits{MaxReadBytes: 1024},
	}
	callerCtx, cancelCaller := context.WithCancel(context.Background())
	cancelCaller()
	execCtx, cancelExec := context.WithCancel(context.Background())
	defer cancelExec()
	h := &singleRequestHandle{
		req:             SingleRequestRequest{RequestID: "request-cancelled"},
		binding:         binding,
		state:           SingleRequestStatePlanning,
		lastSequence:    1,
		progressCh:      make(chan SingleRequestProgress, 4),
		doneCh:          make(chan struct{}),
		callerCtx:       callerCtx,
		execCtx:         execCtx,
		cancelExec:      cancelExec,
		requestDeadline: now.Add(time.Second),
		cleanupComplete: true,
		toolLoop: singleRequestToolLoopState{
			continuation:  newScriptedInternalToolExecutor(),
			runtime:       &Service{},
			seenCallIDs:   make(map[string]struct{}),
			usage:         make(map[string]singleRequestToolUsage),
			stageDeadline: now.Add(-time.Second),
		},
		timing: newSingleRequestTimingAccumulator(nil, nil),
	}

	err := h.SubmitEnvelope(SingleRequestEnvelope{
		RequestID:  h.req.RequestID,
		Sequence:   2,
		Stage:      SingleRequestStateInternalTool,
		SavedStage: SingleRequestStatePlanning,
		ToolCall: &InternalWorkspaceToolCall{
			RequestID: h.req.RequestID, StageID: "plan", ToolCallID: "tool-cancelled",
			Name: InternalWorkspaceToolRead, Arguments: json.RawMessage(`{"relative_path":"README.md"}`),
		},
	})
	if !errors.Is(err, ErrSingleRequestCancelled) {
		t.Fatalf("SubmitEnvelope error = %v, want cancelled", err)
	}
	if h.State() != SingleRequestStateCancelled {
		t.Fatalf("state = %s, want cancelled", h.State())
	}
	if h.toolWork != 0 || h.toolLoop.pendingCallID != "" || len(h.toolLoop.seenCallIDs) != 0 {
		t.Fatalf("late admission dispatched tool work=%d pending=%q seen=%d", h.toolWork, h.toolLoop.pendingCallID, len(h.toolLoop.seenCallIDs))
	}

	terminalCount := 0
	for progress := range h.progressCh {
		if progress.Terminal != nil {
			terminalCount++
			if progress.Terminal.Kind != SingleRequestTerminalCancelled {
				t.Fatalf("terminal = %+v, want cancelled", progress.Terminal)
			}
		}
	}
	if terminalCount != 1 {
		t.Fatalf("terminal count = %d, want 1", terminalCount)
	}
}

func TestSingleRequestInternalToolLoopStageDeadline(t *testing.T) {
	executor := newScriptedInternalToolExecutor(InternalWorkspaceToolCall{
		ToolCallID: "tool-command", Name: InternalWorkspaceToolCommand,
		Arguments: json.RawMessage(`{"command_id":"test"}`),
	})
	service, node := newInternalToolLoopService(t, executor)
	var openCount atomic.Int32
	installInternalLoopOpenResponder(node, &openCount)
	var sequence atomic.Int32
	toolEntered := make(chan struct{})
	release := make(chan struct{})
	serveWorkspaceConcurrent(&node.Communicator, &sequence, func(req *iop.WorkspaceToolRequest) *iop.WorkspaceToolResponse {
		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"}
	})

	handle, err := service.StartSingleRequest(context.Background(), internalLoopRequest(t, func(binding *SingleRequestBinding) {
		binding.Limits.StageTimeoutMS = 50
	}))
	if err != nil {
		t.Fatalf("StartSingleRequest: %v", err)
	}
	select {
	case <-toolEntered:
	case <-time.After(2 * time.Second):
		close(release)
		t.Fatal("tool did not reach Node")
	}
	if _, err := waitForExecution(t, handle); !errors.Is(err, ErrSingleRequestInternalToolBudget) {
		close(release)
		t.Fatalf("Wait error = %v, want budget exhaustion", err)
	}
	close(release)
	if handle.State() != SingleRequestStateFailed || executor.continueCount.Load() != 0 {
		t.Fatalf("state=%s continuations=%d, want failed/0", handle.State(), executor.continueCount.Load())
	}
}

func TestPrepareInternalWorkspaceToolDeadlineOwnership(t *testing.T) {
	for _, test := range []struct {
		name                string
		requestDeadlineFrom time.Duration
		stageDeadlineFrom   time.Duration
		wantErrorClass      singleRequestErrorClass
	}{
		{
			name:                "request deadline wins after both deadlines expire",
			requestDeadlineFrom: -2 * time.Second,
			stageDeadlineFrom:   -time.Second,
			wantErrorClass:      singleRequestErrorClassInternalToolBudget,
		},
		{
			name:                "earlier stage deadline remains timeout",
			requestDeadlineFrom: time.Second,
			stageDeadlineFrom:   -time.Second,
			wantErrorClass:      singleRequestErrorClassTimeout,
		},
	} {
		t.Run(test.name, func(t *testing.T) {
			now := time.Now()
			binding := createTestBinding(t)
			binding.Workspace = &SingleRequestWorkspaceBinding{
				Ref:          "workspace-loop",
				OperationIDs: []string{"read"},
				Limits:       SingleRequestWorkspaceLimits{MaxReadBytes: 1024},
			}
			h := &singleRequestHandle{
				req:             SingleRequestRequest{RequestID: "request-deadline"},
				binding:         binding,
				state:           SingleRequestStatePlanning,
				callerCtx:       context.Background(),
				execCtx:         context.Background(),
				requestDeadline: now.Add(test.requestDeadlineFrom),
				toolLoop: singleRequestToolLoopState{
					continuation:  newScriptedInternalToolExecutor(),
					runtime:       &Service{},
					seenCallIDs:   make(map[string]struct{}),
					usage:         make(map[string]singleRequestToolUsage),
					stageDeadline: now.Add(test.stageDeadlineFrom),
				},
			}

			pending, err, errorClass := h.prepareInternalWorkspaceToolLocked(&InternalWorkspaceToolCall{
				RequestID: "request-deadline", StageID: "plan", ToolCallID: "tool-deadline",
				Name: InternalWorkspaceToolRead, Arguments: json.RawMessage(`{"relative_path":"README.md"}`),
			})
			if pending != nil || !errors.Is(err, ErrSingleRequestInternalToolBudget) || errorClass != test.wantErrorClass {
				t.Fatalf("prepare = (%+v, %v, %q), want (nil, internal tool budget, %q)", pending, err, errorClass, test.wantErrorClass)
			}
		})
	}
}

func TestSingleRequestInternalToolRequestWallClockBudgetOwnership(t *testing.T) {
	const (
		iterations  = 20
		rawSentinel = "RAW-TOOL-BUDGET-SENTINEL"
	)
	for iteration := 0; iteration < iterations; iteration++ {
		executor := newScriptedInternalToolExecutor(InternalWorkspaceToolCall{
			ToolCallID: "request-budget-tool", Name: InternalWorkspaceToolRead,
			Arguments: json.RawMessage(`{"relative_path":"` + rawSentinel + `.txt"}`),
		})
		service, node := newInternalToolLoopService(t, executor)
		observer := &capturingObserver{}
		service.SetSingleRequestObserver(observer)
		var openCount atomic.Int32
		installInternalLoopOpenResponder(node, &openCount)
		var sequence atomic.Int32
		toolEntered := make(chan struct{})
		release := make(chan struct{})
		serveWorkspaceConcurrent(&node.Communicator, &sequence, func(req *iop.WorkspaceToolRequest) *iop.WorkspaceToolResponse {
			close(toolEntered)
			<-release
			return &iop.WorkspaceToolResponse{
				RequestId: req.GetRequestId(), StageId: req.GetStageId(), ToolCallId: req.GetToolCallId(),
				Status: iop.WorkspaceStatus_WORKSPACE_STATUS_TIMEOUT, ErrorCode: iop.WorkspaceErrorCode_WORKSPACE_ERROR_CODE_TIMEOUT,
			}
		})
		serveWorkspaceConcurrent(&node.Communicator, &sequence, func(req *iop.WorkspaceCancelRequest) *iop.WorkspaceCancelResponse {
			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,
			}
		})

		handle, err := service.StartSingleRequest(context.Background(), internalLoopRequest(t, func(binding *SingleRequestBinding) {
			binding.Limits.WallClockMS = 30
			binding.Limits.StageTimeoutMS = 30
		}))
		if err != nil {
			close(release)
			t.Fatalf("iteration=%d StartSingleRequest: %v", iteration, err)
		}
		select {
		case <-toolEntered:
		case <-time.After(2 * time.Second):
			close(release)
			t.Fatalf("iteration=%d tool did not reach Node", iteration)
		}
		internal := handle.(*singleRequestHandle)
		select {
		case <-internal.execCtx.Done():
		case <-time.After(2 * time.Second):
			close(release)
			t.Fatalf("iteration=%d request wall-clock did not expire", iteration)
		}
		close(release)
		assertSingleRequestRequestBudgetOwnership(t, handle, observer, ErrSingleRequestInternalToolBudget, rawSentinel)
		if openCount.Load() != 1 || executor.continueCount.Load() != 0 {
			t.Fatalf("iteration=%d open/continuation = %d/%d, want 1/0", iteration, openCount.Load(), executor.continueCount.Load())
		}
	}
}