DeepSeek-Reasonix/internal/agent/execute_batch.go

package agent

import (
	"context"
	"encoding/json"
	"fmt"
	"sync"
	"time"

	"reasonix/internal/event"
	"reasonix/internal/provider"
	"reasonix/internal/tool"
)

// toolOutcome is one tool call's result. output is the first-visible bounded
// form the model sees; rawOutput is the full original when truncation applied
// (empty when identical so we avoid double storage). images ride outside text.
type toolOutcome struct {
	runState                   provider.ToolRunState
	visionSummary              *provider.VisionSummary
	output                     string
	rawOutput                  string // full original when different from output
	images                     []string
	blocked                    bool
	errMsg                     string
	truncated                  bool
	truncMsg                   string
	resolved                   bool
	resolvedName               string
	capabilityID               string
	resolvedReadOnly, executed bool
	workspaceMutation          *event.WorkspaceMutation
	effective                  workspaceEffectiveCall
	// execution is local shell metadata (optional). Provider messages strip it
	// via ModelMessages; UI/event sinks surface it on ToolResult cards.
	execution *tool.ShellExecution
	// mcpApp is the optional MCP Apps presentation; provider-excluded like
	// execution, persisted for Desktop cards.
	mcpApp *provider.MCPAppPresentation
	// presentedFiles is trusted host metadata from a successful built-in
	// present call. It shares the persisted tool-result commit boundary.
	presentedFiles   []provider.PresentedFile
	readTaskID       string
	readEnvelope     *tool.ReadResultEnvelope
	diagnostic       *tool.OperationDiagnostic
	evidenceSource   tool.EvidenceTargetInfo
	readActiveMillis int64
	subagentOutcome  *SubagentOutcome
}

// batchExecution is the result of one provider tool-call batch.
type batchExecution struct {
	results    []string
	outcomes   []toolOutcome
	images     [][]string
	executions []*tool.ShellExecution
	err        error
}

// executeBatch dispatches one model turn's tool calls. ToolDispatch events are
// emitted up front in call order; contiguous known ReadOnly calls fan out
// across goroutines while unknown and writer calls run serially so write/read
// ordering stays provider-ordered. Each completed serial call (or read-only
// group) is checkpointed before the next group starts.
func (a *Agent) executeBatch(ctx context.Context, turn *turnRuntime, calls []provider.ToolCall) batchExecution {
	// The assistant message already stored this slice in Session. Keep execution
	// state separate so refreshing a dependent preview never mutates shared
	// session memory outside Session's lock.
	calls = append([]provider.ToolCall(nil), calls...)
	if err := a.prepareToolBatch(ctx, calls); err != nil {
		return batchExecution{err: err}
	}

	slots := newBatchSlots(calls)
	results, outcomes, durations, startedAt := slots.results, slots.outcomes, slots.durations, slots.startedAt
	ranParallel := make([]bool, len(calls))
	batchStart := time.Now()
	// Full dispatches used the batch's initial file state. After a writer runs
	// (even a failed one — disk may have mutated), refresh dependent writer
	// previews. The first writer stays on the single-preview fast path.
	earlierWriterRan := false
	surfaceWriters := slots.surfaceWriters
	var batchErr error
	var batchErrOnce sync.Once
	run := func(s *batchSlots, i int) {
		t, _, ambiguous := a.svc.tools.ResolveCall(s.calls[i].Name)
		known := t != nil && len(ambiguous) == 0
		writer := known && !t.ReadOnly()
		s.surfaceWriters[i] = writer
		if earlierWriterRan && writer {
			if refreshed, changed := refreshCurrentFileDiff(ctx, t, s.calls[i]); changed {
				s.calls[i] = refreshed
				a.sess.conversation.UpdateToolCallPreview(refreshed)
				if err := a.emitFullToolDispatch(ctx, refreshed, true); err != nil {
					wrapped := fmt.Errorf("persist refreshed tool dispatch %s: %w", refreshed.ID, err)
					batchErrOnce.Do(func() { batchErr = wrapped })
					s.outcomes[i] = toolOutcome{output: "cancelled: tool dispatch was not durable", errMsg: wrapped.Error()}
					s.results[i] = s.outcomes[i].output
					return
				}
			}
		}
		start := time.Now()
		s.startedAt[i] = start.UnixMilli()
		s.outcomes[i] = a.executeOne(ctx, turn, s.calls[i])
		recordWorkspaceMutation(a.svc.sink, s.outcomes[i].workspaceMutation)
		if s.outcomes[i].executed {
			s.surfaceWriters[i] = s.outcomes[i].workspaceMutation != nil
		}
		if s.outcomes[i].resolved {
			readOnly := s.outcomes[i].resolvedReadOnly
			s.calls[i].ResolvedName = s.outcomes[i].resolvedName
			s.calls[i].CapabilityID = s.outcomes[i].capabilityID
			s.calls[i].ResolvedReadOnly = &readOnly
			s.surfaceWriters[i] = !readOnly
		}
		s.durations[i] = time.Since(start).Milliseconds()
		s.results[i] = s.outcomes[i].output
	}
	committed := make([]bool, len(calls))
	committedMessages := make([]provider.Message, len(calls))
	finalize := func(i int) {
		if committed[i] {
			return
		}
		committed[i] = true
		results[i] = outcomes[i].output
		oneResult := results[i : i+1]
		a.applyRepeatReminders(calls[i:i+1], oneResult)
		outcomes[i].output = results[i]
		a.commitBatchCallResolution(ctx, calls[i])
		a.finishToolRecovery(calls[i], outcomes[i])
		committedMessage := a.buildBatchToolResult(ctx, calls[i], outcomes[i])
		committedMessages[i] = committedMessage
		if err := a.emitBatchToolResult(ctx, calls[i], outcomes[i], committedMessage, durations[i], startedAt[i], ranParallel[i], batchStart); err != nil {
			batchErrOnce.Do(func() { batchErr = fmt.Errorf("persist tool result %s: %w", calls[i].ID, err) })
		} else {
			a.sess.conversation.Add(committedMessage)
		}
		if surfaceWriters[i] || (outcomes[i].resolved && !outcomes[i].resolvedReadOnly) {
			earlierWriterRan = true
		}
	}
	cancelled := false
	markCancelled := func(start int) {
		errMsg := context.Canceled.Error()
		if err := ctx.Err(); err != nil {
			errMsg = err.Error()
		}
		output := "cancelled: context cancelled before execution"
		for j := start; j < len(calls); j++ {
			results[j] = output
			outcomes[j] = toolOutcome{output: output, errMsg: errMsg}
		}
		cancelled = true
	}

	for _, batch := range a.toolCallBatches(calls) {
		if ctx.Err() != nil || batchErr != nil {
			markCancelled(batch.start)
			break
		}
		if batch.parallel && batch.end-batch.start > 1 {
			// Parallel segments are read-only by construction; no mutation barrier.
			private := slots.fork()
			ranUntil, finished := runParallel(ctx, batch.start, batch.end, func(i int) {
				a.stragglers.enter()
				defer a.stragglers.leave()
				run(private, i)
			})
			for i := batch.start; i < ranUntil; i++ {
				if finished[i] {
					slots.adopt(private, i)
				} else {
					slots.abandon(i)
				}
				ranParallel[i] = true
				finalize(i)
			}
			// After parallel execution completes, check if context was cancelled.
			// The individual tool executions should have detected ctx.Done(), but
			// we verify here to ensure we don't continue to subsequent batches.
			if ctx.Err() != nil {
				markCancelled(ranUntil)
				break
			}
			continue
		}
		for i := batch.start; i < batch.end; i++ {
			// Before executing the next tool, check if context was cancelled.
			// This prevents starting new tools when a previous tool's execution
			// triggered cancellation.
			if ctx.Err() != nil || batchErr != nil {
				markCancelled(i)
				break
			}
			run(slots, i)
			finalize(i)
			// After each tool execution, also check if the context was cancelled.
			// If so, stop executing remaining tools and return immediately so
			// the agent loop can detect the cancellation and exit.
			if ctx.Err() != nil {
				markCancelled(i + 1)
				break
			}
		}
		if cancelled {
			break
		}
	}

	for i := range calls {
		finalize(i)
	}
	return completeBatchExecution(ctx, calls, results, outcomes, committedMessages, committed, batchErr)
}

func completeBatchExecution(ctx context.Context, calls []provider.ToolCall, results []string, outcomes []toolOutcome, committedMessages []provider.Message, committed []bool, batchErr error) batchExecution {
	if err := validateBatchToolResultCorrespondence(calls, committedMessages, committed); err != nil && batchErr == nil {
		batchErr = err
	}
	images := make([][]string, len(calls))
	executions := make([]*tool.ShellExecution, len(calls))
	for i := range outcomes {
		images[i] = outcomes[i].images
		executions[i] = outcomes[i].execution
	}
	if batchErr == nil && ctx.Err() != nil {
		batchErr = ctx.Err()
	}
	return batchExecution{
		results:    results,
		outcomes:   outcomes,
		images:     images,
		executions: executions,
		err:        batchErr,
	}
}

// validateBatchToolResultCorrespondence is the provider boundary invariant:
// every executed call yields exactly one result at the same index and with the
// same call id. Equal result bodies are intentionally irrelevant.
func validateBatchToolResultCorrespondence(calls []provider.ToolCall, results []provider.Message, committed []bool) error {
	if len(results) != len(calls) || len(committed) != len(calls) {
		return fmt.Errorf("tool result correspondence: %d calls, %d results, %d commit markers", len(calls), len(results), len(committed))
	}
	for i := range calls {
		if !committed[i] {
			return fmt.Errorf("tool result correspondence: call %q at index %d has no result", calls[i].ID, i)
		}
		result := results[i]
		if result.Role != provider.RoleTool || result.ToolCallID != calls[i].ID {
			return fmt.Errorf("tool result correspondence: call %q at index %d got role=%q call_id=%q", calls[i].ID, i, result.Role, result.ToolCallID)
		}
	}
	return nil
}

func (a *Agent) commitBatchCallResolution(ctx context.Context, call provider.ToolCall) {
	if call.ResolvedReadOnly == nil {
		return
	}
	a.sess.conversation.UpdateToolCallResolution(call)
	a.emitResolvedToolDispatch(ctx, call)
}

type toolCallBatch struct {
	start    int
	end      int
	parallel bool
}

// toolCallBatches preserves read-only fan-out unless a tool hook can mutate the
// workspace. Such hooks are covered by a whole-workspace claim, so their calls
// must run in provider order instead of racing that claim against each other.
func (a *Agent) toolCallBatches(calls []provider.ToolCall) []toolCallBatch {
	batches := partitionToolCalls(a.svc.tools, calls)
	if !toolHooksMayMutateWorkspace(a.svc.hooks) {
		return batches
	}
	for i := range batches {
		batches[i].parallel = false
	}
	return batches
}

// partitionToolCalls keeps provider order while letting contiguous known
// read-only tools run together; unknown and writer tools are single-call
// serial batches. State tools stay serial so provider order stays result
// order; use_capability is serial as it may resolve to a real MCP writer.
func partitionToolCalls(r *tool.Registry, calls []provider.ToolCall) []toolCallBatch {
	var batches []toolCallBatch
	for i := 0; i < len(calls); {
		if parallelisableCall(r, calls[i]) {
			start := i
			i++
			for i < len(calls) && parallelisableCall(r, calls[i]) {
				i++
			}
			batches = append(batches, toolCallBatch{start: start, end: i, parallel: true})
			continue
		}
		batches = append(batches, toolCallBatch{start: i, end: i + 1})
		i++
	}
	return batches
}

func parallelisableCall(r *tool.Registry, call provider.ToolCall) bool {
	switch call.Name {
	case "todo_write", "get_goal", "create_goal", "update_goal", "job_output", "wait", "bash_output", "compress":
		return false
	}
	target, _, ambiguous := r.ResolveCall(call.Name)
	if target == nil || len(ambiguous) != 0 {
		return false
	}
	if classifier, ok := target.(tool.BatchClassifier); ok {
		class := classifier.ClassifyCall(json.RawMessage(call.Arguments))
		return class.Known && class.ReadOnly && class.ParallelSafe
	}
	if _, dynamic := target.(tool.CallResolver); dynamic {
		return false
	}
	return target.ReadOnly()
}

// parallelStragglerGrace bounds how long a cancelled parallel segment waits for
// tools that have not returned. Tool owners kill their own processes within
// their WaitDelay; past this the batch reports the effect as unknown instead
// of keeping the whole turn wedged behind one call that ignores its context.
var parallelStragglerGrace = 15 * time.Second

// runParallel returns the launched prefix and which of those calls finished.
// An unfinished index belongs to a straggler that still owns its private slot.
func runParallel(ctx context.Context, start, end int, run func(int)) (int, []bool) {
	const maxParallel = 8
	sem := make(chan struct{}, maxParallel)
	var wg sync.WaitGroup
	completed := make(chan int, end-start)
	ranUntil := start
launch:
	for i := start; i < end; i++ {
		if ctx.Err() != nil {
			break
		}
		select {
		case sem <- struct{}{}:
		case <-ctx.Done():
			break launch
		}
		if ctx.Err() != nil {
			<-sem
			break
		}

		wg.Add(1)
		ranUntil = i + 1
		go func() {
			defer wg.Done()
			defer func() { <-sem }()
			run(i)
			completed <- i
		}()
	}
	allDone := make(chan struct{})
	go func() {
		wg.Wait()
		close(allDone)
	}()
	select {
	case <-allDone:
	case <-ctx.Done():
		select {
		case <-allDone:
		case <-time.After(parallelStragglerGrace):
		}
	}
	finished := make([]bool, end)
	for {
		select {
		case i := <-completed:
			finished[i] = true
		default:
			return ranUntil, finished
		}
	}
}

// batchSlots is one batch's per-call execution state. Parallel segments run
// against a fork so a tool that outlives cancellation writes only into slots
// the batch has already stopped reading.
type batchSlots struct {
	calls          []provider.ToolCall
	outcomes       []toolOutcome
	results        []string
	durations      []int64
	startedAt      []int64
	surfaceWriters []bool
}

func newBatchSlots(calls []provider.ToolCall) *batchSlots {
	n := len(calls)
	return &batchSlots{
		calls: calls, outcomes: make([]toolOutcome, n), results: make([]string, n),
		durations: make([]int64, n), startedAt: make([]int64, n), surfaceWriters: make([]bool, n),
	}
}

func (s *batchSlots) fork() *batchSlots {
	return newBatchSlots(append([]provider.ToolCall(nil), s.calls...))
}

func (s *batchSlots) adopt(from *batchSlots, i int) {
	s.calls[i], s.outcomes[i], s.results[i] = from.calls[i], from.outcomes[i], from.results[i]
	s.durations[i], s.startedAt[i], s.surfaceWriters[i] = from.durations[i], from.startedAt[i], from.surfaceWriters[i]
}

const abandonedToolOutput = "interrupted: the tool did not stop after cancellation; its effect is unknown"

func (s *batchSlots) abandon(i int) {
	s.outcomes[i] = toolOutcome{output: abandonedToolOutput, errMsg: abandonedToolOutput, executed: true}
	s.results[i] = abandonedToolOutput
}