| 1 | package agent |
| 2 | |
| 3 | import ( |
| 4 | "context" |
| 5 | "encoding/json" |
| 6 | "fmt" |
| 7 | "sync" |
| 8 | "time" |
| 9 | |
| 10 | "reasonix/internal/event" |
| 11 | "reasonix/internal/provider" |
| 12 | "reasonix/internal/tool" |
| 13 | ) |
| 14 | |
| 15 | // toolOutcome is one tool call's result. output is the first-visible bounded |
| 16 | // form the model sees; rawOutput is the full original when truncation applied |
| 17 | // (empty when identical so we avoid double storage). images ride outside text. |
| 18 | type toolOutcome struct { |
| 19 | runState provider.ToolRunState |
| 20 | visionSummary *provider.VisionSummary |
| 21 | output string |
| 22 | rawOutput string // full original when different from output |
| 23 | images []string |
| 24 | blocked bool |
| 25 | errMsg string |
| 26 | truncated bool |
| 27 | truncMsg string |
| 28 | resolved bool |
| 29 | resolvedName string |
| 30 | capabilityID string |
| 31 | resolvedReadOnly, executed bool |
| 32 | workspaceMutation *event.WorkspaceMutation |
| 33 | effective workspaceEffectiveCall |
| 34 | // execution is local shell metadata (optional). Provider messages strip it |
| 35 | // via ModelMessages; UI/event sinks surface it on ToolResult cards. |
| 36 | execution *tool.ShellExecution |
| 37 | // mcpApp is the optional MCP Apps presentation; provider-excluded like |
| 38 | // execution, persisted for Desktop cards. |
| 39 | mcpApp *provider.MCPAppPresentation |
| 40 | // presentedFiles is trusted host metadata from a successful built-in |
| 41 | // present call. It shares the persisted tool-result commit boundary. |
| 42 | presentedFiles []provider.PresentedFile |
| 43 | readTaskID string |
| 44 | readEnvelope *tool.ReadResultEnvelope |
| 45 | diagnostic *tool.OperationDiagnostic |
| 46 | evidenceSource tool.EvidenceTargetInfo |
| 47 | readActiveMillis int64 |
| 48 | subagentOutcome *SubagentOutcome |
| 49 | } |
| 50 | |
| 51 | // batchExecution is the result of one provider tool-call batch. |
| 52 | type batchExecution struct { |
| 53 | results []string |
| 54 | outcomes []toolOutcome |
| 55 | images [][]string |
| 56 | executions []*tool.ShellExecution |
| 57 | err error |
| 58 | } |
| 59 | |
| 60 | // executeBatch dispatches one model turn's tool calls. ToolDispatch events are |
| 61 | // emitted up front in call order; contiguous known ReadOnly calls fan out |
| 62 | // across goroutines while unknown and writer calls run serially so write/read |
| 63 | // ordering stays provider-ordered. Each completed serial call (or read-only |
| 64 | // group) is checkpointed before the next group starts. |
| 65 | func (a *Agent) executeBatch(ctx context.Context, turn *turnRuntime, calls []provider.ToolCall) batchExecution { |
| 66 | // The assistant message already stored this slice in Session. Keep execution |
| 67 | // state separate so refreshing a dependent preview never mutates shared |
| 68 | // session memory outside Session's lock. |
| 69 | calls = append([]provider.ToolCall(nil), calls...) |
| 70 | if err := a.prepareToolBatch(ctx, calls); err != nil { |
| 71 | return batchExecution{err: err} |
| 72 | } |
| 73 | |
| 74 | slots := newBatchSlots(calls) |
| 75 | results, outcomes, durations, startedAt := slots.results, slots.outcomes, slots.durations, slots.startedAt |
| 76 | ranParallel := make([]bool, len(calls)) |
| 77 | batchStart := time.Now() |
| 78 | // Full dispatches used the batch's initial file state. After a writer runs |
| 79 | // (even a failed one — disk may have mutated), refresh dependent writer |
| 80 | // previews. The first writer stays on the single-preview fast path. |
| 81 | earlierWriterRan := false |
| 82 | surfaceWriters := slots.surfaceWriters |
| 83 | var batchErr error |
| 84 | var batchErrOnce sync.Once |
| 85 | run := func(s *batchSlots, i int) { |
| 86 | t, _, ambiguous := a.svc.tools.ResolveCall(s.calls[i].Name) |
| 87 | known := t != nil && len(ambiguous) == 0 |
| 88 | writer := known && !t.ReadOnly() |
| 89 | s.surfaceWriters[i] = writer |
| 90 | if earlierWriterRan && writer { |
| 91 | if refreshed, changed := refreshCurrentFileDiff(ctx, t, s.calls[i]); changed { |
| 92 | s.calls[i] = refreshed |
| 93 | a.sess.conversation.UpdateToolCallPreview(refreshed) |
| 94 | if err := a.emitFullToolDispatch(ctx, refreshed, true); err != nil { |
| 95 | wrapped := fmt.Errorf("persist refreshed tool dispatch %s: %w", refreshed.ID, err) |
| 96 | batchErrOnce.Do(func() { batchErr = wrapped }) |
| 97 | s.outcomes[i] = toolOutcome{output: "cancelled: tool dispatch was not durable", errMsg: wrapped.Error()} |
| 98 | s.results[i] = s.outcomes[i].output |
| 99 | return |
| 100 | } |
| 101 | } |
| 102 | } |
| 103 | start := time.Now() |
| 104 | s.startedAt[i] = start.UnixMilli() |
| 105 | s.outcomes[i] = a.executeOne(ctx, turn, s.calls[i]) |
| 106 | recordWorkspaceMutation(a.svc.sink, s.outcomes[i].workspaceMutation) |
| 107 | if s.outcomes[i].executed { |
| 108 | s.surfaceWriters[i] = s.outcomes[i].workspaceMutation != nil |
| 109 | } |
| 110 | if s.outcomes[i].resolved { |
| 111 | readOnly := s.outcomes[i].resolvedReadOnly |
| 112 | s.calls[i].ResolvedName = s.outcomes[i].resolvedName |
| 113 | s.calls[i].CapabilityID = s.outcomes[i].capabilityID |
| 114 | s.calls[i].ResolvedReadOnly = &readOnly |
| 115 | s.surfaceWriters[i] = !readOnly |
| 116 | } |
| 117 | s.durations[i] = time.Since(start).Milliseconds() |
| 118 | s.results[i] = s.outcomes[i].output |
| 119 | } |
| 120 | committed := make([]bool, len(calls)) |
| 121 | committedMessages := make([]provider.Message, len(calls)) |
| 122 | finalize := func(i int) { |
| 123 | if committed[i] { |
| 124 | return |
| 125 | } |
| 126 | committed[i] = true |
| 127 | results[i] = outcomes[i].output |
| 128 | oneResult := results[i : i+1] |
| 129 | a.applyRepeatReminders(calls[i:i+1], oneResult) |
| 130 | outcomes[i].output = results[i] |
| 131 | a.commitBatchCallResolution(ctx, calls[i]) |
| 132 | a.finishToolRecovery(calls[i], outcomes[i]) |
| 133 | committedMessage := a.buildBatchToolResult(ctx, calls[i], outcomes[i]) |
| 134 | committedMessages[i] = committedMessage |
| 135 | if err := a.emitBatchToolResult(ctx, calls[i], outcomes[i], committedMessage, durations[i], startedAt[i], ranParallel[i], batchStart); err != nil { |
| 136 | batchErrOnce.Do(func() { batchErr = fmt.Errorf("persist tool result %s: %w", calls[i].ID, err) }) |
| 137 | } else { |
| 138 | a.sess.conversation.Add(committedMessage) |
| 139 | } |
| 140 | if surfaceWriters[i] || (outcomes[i].resolved && !outcomes[i].resolvedReadOnly) { |
| 141 | earlierWriterRan = true |
| 142 | } |
| 143 | } |
| 144 | cancelled := false |
| 145 | markCancelled := func(start int) { |
| 146 | errMsg := context.Canceled.Error() |
| 147 | if err := ctx.Err(); err != nil { |
| 148 | errMsg = err.Error() |
| 149 | } |
| 150 | output := "cancelled: context cancelled before execution" |
| 151 | for j := start; j < len(calls); j++ { |
| 152 | results[j] = output |
| 153 | outcomes[j] = toolOutcome{output: output, errMsg: errMsg} |
| 154 | } |
| 155 | cancelled = true |
| 156 | } |
| 157 | |
| 158 | for _, batch := range a.toolCallBatches(calls) { |
| 159 | if ctx.Err() != nil || batchErr != nil { |
| 160 | markCancelled(batch.start) |
| 161 | break |
| 162 | } |
| 163 | if batch.parallel && batch.end-batch.start > 1 { |
| 164 | // Parallel segments are read-only by construction; no mutation barrier. |
| 165 | private := slots.fork() |
| 166 | ranUntil, finished := runParallel(ctx, batch.start, batch.end, func(i int) { |
| 167 | a.stragglers.enter() |
| 168 | defer a.stragglers.leave() |
| 169 | run(private, i) |
| 170 | }) |
| 171 | for i := batch.start; i < ranUntil; i++ { |
| 172 | if finished[i] { |
| 173 | slots.adopt(private, i) |
| 174 | } else { |
| 175 | slots.abandon(i) |
| 176 | } |
| 177 | ranParallel[i] = true |
| 178 | finalize(i) |
| 179 | } |
| 180 | // After parallel execution completes, check if context was cancelled. |
| 181 | // The individual tool executions should have detected ctx.Done(), but |
| 182 | // we verify here to ensure we don't continue to subsequent batches. |
| 183 | if ctx.Err() != nil { |
| 184 | markCancelled(ranUntil) |
| 185 | break |
| 186 | } |
| 187 | continue |
| 188 | } |
| 189 | for i := batch.start; i < batch.end; i++ { |
| 190 | // Before executing the next tool, check if context was cancelled. |
| 191 | // This prevents starting new tools when a previous tool's execution |
| 192 | // triggered cancellation. |
| 193 | if ctx.Err() != nil || batchErr != nil { |
| 194 | markCancelled(i) |
| 195 | break |
| 196 | } |
| 197 | run(slots, i) |
| 198 | finalize(i) |
| 199 | // After each tool execution, also check if the context was cancelled. |
| 200 | // If so, stop executing remaining tools and return immediately so |
| 201 | // the agent loop can detect the cancellation and exit. |
| 202 | if ctx.Err() != nil { |
| 203 | markCancelled(i + 1) |
| 204 | break |
| 205 | } |
| 206 | } |
| 207 | if cancelled { |
| 208 | break |
| 209 | } |
| 210 | } |
| 211 | |
| 212 | for i := range calls { |
| 213 | finalize(i) |
| 214 | } |
| 215 | return completeBatchExecution(ctx, calls, results, outcomes, committedMessages, committed, batchErr) |
| 216 | } |
| 217 | |
| 218 | func completeBatchExecution(ctx context.Context, calls []provider.ToolCall, results []string, outcomes []toolOutcome, committedMessages []provider.Message, committed []bool, batchErr error) batchExecution { |
| 219 | if err := validateBatchToolResultCorrespondence(calls, committedMessages, committed); err != nil && batchErr == nil { |
| 220 | batchErr = err |
| 221 | } |
| 222 | images := make([][]string, len(calls)) |
| 223 | executions := make([]*tool.ShellExecution, len(calls)) |
| 224 | for i := range outcomes { |
| 225 | images[i] = outcomes[i].images |
| 226 | executions[i] = outcomes[i].execution |
| 227 | } |
| 228 | if batchErr == nil && ctx.Err() != nil { |
| 229 | batchErr = ctx.Err() |
| 230 | } |
| 231 | return batchExecution{ |
| 232 | results: results, |
| 233 | outcomes: outcomes, |
| 234 | images: images, |
| 235 | executions: executions, |
| 236 | err: batchErr, |
| 237 | } |
| 238 | } |
| 239 | |
| 240 | // validateBatchToolResultCorrespondence is the provider boundary invariant: |
| 241 | // every executed call yields exactly one result at the same index and with the |
| 242 | // same call id. Equal result bodies are intentionally irrelevant. |
| 243 | func validateBatchToolResultCorrespondence(calls []provider.ToolCall, results []provider.Message, committed []bool) error { |
| 244 | if len(results) != len(calls) || len(committed) != len(calls) { |
| 245 | return fmt.Errorf("tool result correspondence: %d calls, %d results, %d commit markers", len(calls), len(results), len(committed)) |
| 246 | } |
| 247 | for i := range calls { |
| 248 | if !committed[i] { |
| 249 | return fmt.Errorf("tool result correspondence: call %q at index %d has no result", calls[i].ID, i) |
| 250 | } |
| 251 | result := results[i] |
| 252 | if result.Role != provider.RoleTool || result.ToolCallID != calls[i].ID { |
| 253 | 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) |
| 254 | } |
| 255 | } |
| 256 | return nil |
| 257 | } |
| 258 | |
| 259 | func (a *Agent) commitBatchCallResolution(ctx context.Context, call provider.ToolCall) { |
| 260 | if call.ResolvedReadOnly == nil { |
| 261 | return |
| 262 | } |
| 263 | a.sess.conversation.UpdateToolCallResolution(call) |
| 264 | a.emitResolvedToolDispatch(ctx, call) |
| 265 | } |
| 266 | |
| 267 | type toolCallBatch struct { |
| 268 | start int |
| 269 | end int |
| 270 | parallel bool |
| 271 | } |
| 272 | |
| 273 | // toolCallBatches preserves read-only fan-out unless a tool hook can mutate the |
| 274 | // workspace. Such hooks are covered by a whole-workspace claim, so their calls |
| 275 | // must run in provider order instead of racing that claim against each other. |
| 276 | func (a *Agent) toolCallBatches(calls []provider.ToolCall) []toolCallBatch { |
| 277 | batches := partitionToolCalls(a.svc.tools, calls) |
| 278 | if !toolHooksMayMutateWorkspace(a.svc.hooks) { |
| 279 | return batches |
| 280 | } |
| 281 | for i := range batches { |
| 282 | batches[i].parallel = false |
| 283 | } |
| 284 | return batches |
| 285 | } |
| 286 | |
| 287 | // partitionToolCalls keeps provider order while letting contiguous known |
| 288 | // read-only tools run together; unknown and writer tools are single-call |
| 289 | // serial batches. State tools stay serial so provider order stays result |
| 290 | // order; use_capability is serial as it may resolve to a real MCP writer. |
| 291 | func partitionToolCalls(r *tool.Registry, calls []provider.ToolCall) []toolCallBatch { |
| 292 | var batches []toolCallBatch |
| 293 | for i := 0; i < len(calls); { |
| 294 | if parallelisableCall(r, calls[i]) { |
| 295 | start := i |
| 296 | i++ |
| 297 | for i < len(calls) && parallelisableCall(r, calls[i]) { |
| 298 | i++ |
| 299 | } |
| 300 | batches = append(batches, toolCallBatch{start: start, end: i, parallel: true}) |
| 301 | continue |
| 302 | } |
| 303 | batches = append(batches, toolCallBatch{start: i, end: i + 1}) |
| 304 | i++ |
| 305 | } |
| 306 | return batches |
| 307 | } |
| 308 | |
| 309 | func parallelisableCall(r *tool.Registry, call provider.ToolCall) bool { |
| 310 | switch call.Name { |
| 311 | case "todo_write", "get_goal", "create_goal", "update_goal", "job_output", "wait", "bash_output", "compress": |
| 312 | return false |
| 313 | } |
| 314 | target, _, ambiguous := r.ResolveCall(call.Name) |
| 315 | if target == nil || len(ambiguous) != 0 { |
| 316 | return false |
| 317 | } |
| 318 | if classifier, ok := target.(tool.BatchClassifier); ok { |
| 319 | class := classifier.ClassifyCall(json.RawMessage(call.Arguments)) |
| 320 | return class.Known && class.ReadOnly && class.ParallelSafe |
| 321 | } |
| 322 | if _, dynamic := target.(tool.CallResolver); dynamic { |
| 323 | return false |
| 324 | } |
| 325 | return target.ReadOnly() |
| 326 | } |
| 327 | |
| 328 | // parallelStragglerGrace bounds how long a cancelled parallel segment waits for |
| 329 | // tools that have not returned. Tool owners kill their own processes within |
| 330 | // their WaitDelay; past this the batch reports the effect as unknown instead |
| 331 | // of keeping the whole turn wedged behind one call that ignores its context. |
| 332 | var parallelStragglerGrace = 15 * time.Second |
| 333 | |
| 334 | // runParallel returns the launched prefix and which of those calls finished. |
| 335 | // An unfinished index belongs to a straggler that still owns its private slot. |
| 336 | func runParallel(ctx context.Context, start, end int, run func(int)) (int, []bool) { |
| 337 | const maxParallel = 8 |
| 338 | sem := make(chan struct{}, maxParallel) |
| 339 | var wg sync.WaitGroup |
| 340 | completed := make(chan int, end-start) |
| 341 | ranUntil := start |
| 342 | launch: |
| 343 | for i := start; i < end; i++ { |
| 344 | if ctx.Err() != nil { |
| 345 | break |
| 346 | } |
| 347 | select { |
| 348 | case sem <- struct{}{}: |
| 349 | case <-ctx.Done(): |
| 350 | break launch |
| 351 | } |
| 352 | if ctx.Err() != nil { |
| 353 | <-sem |
| 354 | break |
| 355 | } |
| 356 | |
| 357 | wg.Add(1) |
| 358 | ranUntil = i + 1 |
| 359 | go func() { |
| 360 | defer wg.Done() |
| 361 | defer func() { <-sem }() |
| 362 | run(i) |
| 363 | completed <- i |
| 364 | }() |
| 365 | } |
| 366 | allDone := make(chan struct{}) |
| 367 | go func() { |
| 368 | wg.Wait() |
| 369 | close(allDone) |
| 370 | }() |
| 371 | select { |
| 372 | case <-allDone: |
| 373 | case <-ctx.Done(): |
| 374 | select { |
| 375 | case <-allDone: |
| 376 | case <-time.After(parallelStragglerGrace): |
| 377 | } |
| 378 | } |
| 379 | finished := make([]bool, end) |
| 380 | for { |
| 381 | select { |
| 382 | case i := <-completed: |
| 383 | finished[i] = true |
| 384 | default: |
| 385 | return ranUntil, finished |
| 386 | } |
| 387 | } |
| 388 | } |
| 389 | |
| 390 | // batchSlots is one batch's per-call execution state. Parallel segments run |
| 391 | // against a fork so a tool that outlives cancellation writes only into slots |
| 392 | // the batch has already stopped reading. |
| 393 | type batchSlots struct { |
| 394 | calls []provider.ToolCall |
| 395 | outcomes []toolOutcome |
| 396 | results []string |
| 397 | durations []int64 |
| 398 | startedAt []int64 |
| 399 | surfaceWriters []bool |
| 400 | } |
| 401 | |
| 402 | func newBatchSlots(calls []provider.ToolCall) *batchSlots { |
| 403 | n := len(calls) |
| 404 | return &batchSlots{ |
| 405 | calls: calls, outcomes: make([]toolOutcome, n), results: make([]string, n), |
| 406 | durations: make([]int64, n), startedAt: make([]int64, n), surfaceWriters: make([]bool, n), |
| 407 | } |
| 408 | } |
| 409 | |
| 410 | func (s *batchSlots) fork() *batchSlots { |
| 411 | return newBatchSlots(append([]provider.ToolCall(nil), s.calls...)) |
| 412 | } |
| 413 | |
| 414 | func (s *batchSlots) adopt(from *batchSlots, i int) { |
| 415 | s.calls[i], s.outcomes[i], s.results[i] = from.calls[i], from.outcomes[i], from.results[i] |
| 416 | s.durations[i], s.startedAt[i], s.surfaceWriters[i] = from.durations[i], from.startedAt[i], from.surfaceWriters[i] |
| 417 | } |
| 418 | |
| 419 | const abandonedToolOutput = "interrupted: the tool did not stop after cancellation; its effect is unknown" |
| 420 | |
| 421 | func (s *batchSlots) abandon(i int) { |
| 422 | s.outcomes[i] = toolOutcome{output: abandonedToolOutput, errMsg: abandonedToolOutput, executed: true} |
| 423 | s.results[i] = abandonedToolOutput |
| 424 | } |
| 425 |