Compare commits

...
Author SHA1 Message Date
KRATOS 7bce9ad9ff perf(agent): concurrent read-only tool batches via capability gate
Upgrade the parallel tool planner from SideEffectRead probes to the PR5
CapabilitiesOf contract: EffectReadOnly + ThreadSafe + auto-allowed, with
resource-key conflict boundaries so same-path reads stay sequential.

Mark audited pure reads ThreadSafe (read_file, read_minified_file,
list_directory, glob, grep, skill, web_fetch). Glob/grep use scopedScan
directory keys (nil for workspace-wide "." so scans do not false-conflict).

Tests cover the capability gate, key conflicts, extendParallelRun windows,
and catalog ThreadSafe audit for concurrent-safe reads.
2026-07-17 17:34:16 +05:30
12 changed files with 286 additions and 28 deletions

No files matched your search

+3 -4
View File
@@ -632,10 +632,9 @@ func Run(ctx context.Context, prompt string, provider Provider, options Options)
// The scan happens lazily at the run's first index — never ahead of a
// pending mutating call — so read-after-write ordering is preserved.
if index >= precomputedEnd {
runEnd := index
for runEnd < len(collected.ToolCalls) && parallelSafeToolCall(registry, collected.ToolCalls[runEnd], options) {
runEnd++
}
// Capability-safe consecutive run (ReadOnly+ThreadSafe, no
// resource-key conflicts). See extendParallelRun.
runEnd := extendParallelRun(registry, collected.ToolCalls, index, options)
if runEnd-index >= 2 {
batchSpan := options.Trace.Span(trace.SpanToolExecution)
precomputed = executeParallelReadBatch(ctx, registry, collected.ToolCalls, index, runEnd, permissionMode, options)
+111 -11
View File
@@ -10,12 +10,23 @@ import (
// Parallel read-ahead for tool batches. When a turn requests several
// independent lookups (read_file + grep + glob is the common shape), executing
// them one after another serializes pure I/O waits. A consecutive run of
// auto-allowed read-only calls is executed concurrently instead; results are
// capability-safe read-only calls is executed concurrently instead; results are
// then consumed in the original call order, so guard counters, message
// ordering, abort semantics, and the surface's call/result event pairing are
// byte-identical to sequential execution. Runs never span a mutating call: a
// read that follows a write must observe the write, so eligibility is decided
// per consecutive run, not per batch.
//
// Eligibility uses the PR5 tool-effect contract (tools.CapabilitiesOf):
//
// Effect == ReadOnly
// AND ThreadSafe == true
// AND auto-allowed (no interactive permission prompt on the hot path)
// AND no resource-key conflict with an earlier call in the same concurrent window
//
// Unknown, mutators, interactive tools, and non-thread-safe reads stay
// sequential. Empty resource keys do not conflict (ThreadSafe is the safety
// gate); shared non-empty keys force a batch boundary.
// maxParallelReadTools bounds concurrent read-only tool executions in a turn.
const maxParallelReadTools = 8
@@ -28,28 +39,117 @@ type precomputedToolResult struct {
}
// parallelSafeToolCall reports whether call may run concurrently with its
// neighbors: the tool must exist, be side-effect-free (SideEffectRead), and be
// auto-allowed for these args, so no interactive prompt or workspace mutation
// is on the hot path. Loop-intercepted tools (ask_user, request_permissions)
// and tool_search (mutates the deferred-tool set) stay sequential.
// neighbors under the PR5 capability contract. Loop-intercepted tools
// (ask_user, request_permissions) and tool_search (mutates the deferred-tool
// set) always stay sequential.
func parallelSafeToolCall(registry *tools.Registry, call ToolCall, options Options) bool {
switch call.Name {
case "ask_user", tools.RequestPermissionsToolName, tools.ToolSearchToolName:
return false
}
tool, found := registry.Get(call.Name)
if !found || tool.Safety().SideEffect != tools.SideEffectRead {
if !found {
return false
}
args := map[string]any{}
if call.Arguments != "" {
if err := decodeToolArguments(call.Arguments, &args); err != nil {
return false
}
caps := tools.CapabilitiesOf(tool)
// Fail-closed: only audited concurrent-safe pure reads.
if caps.Effect != tools.EffectReadOnly || !caps.ThreadSafe {
return false
}
args, ok := decodeCallArgs(call)
if !ok {
return false
}
return effectivePermission(tool, args) == tools.PermissionAllow
}
// decodeCallArgs decodes tool call JSON arguments. Returns false on malformed
// input so the call stays sequential (never panics into the parallel path).
func decodeCallArgs(call ToolCall) (map[string]any, bool) {
args := map[string]any{}
if call.Arguments == "" {
return args, true
}
if err := decodeToolArguments(call.Arguments, &args); err != nil {
return nil, false
}
return args, true
}
// resourceKeysForCall returns the conflict keys for a call, or nil when the
// tool has no ResourceKeys function / no keys for these args.
func resourceKeysForCall(registry *tools.Registry, call ToolCall) []string {
tool, found := registry.Get(call.Name)
if !found {
return nil
}
caps := tools.CapabilitiesOf(tool)
if caps.ResourceKeys == nil {
return nil
}
args, ok := decodeCallArgs(call)
if !ok {
return nil
}
return caps.ResourceKeys(args)
}
// resourceKeysConflict reports whether two key sets share any non-empty key.
// Empty key sets never conflict: ThreadSafe is the eligibility gate; keys only
// refine conflict detection among tools that declare specific resources.
func resourceKeysConflict(a, b []string) bool {
if len(a) == 0 || len(b) == 0 {
return false
}
seen := make(map[string]struct{}, len(a))
for _, key := range a {
if key == "" {
continue
}
seen[key] = struct{}{}
}
for _, key := range b {
if key == "" {
continue
}
if _, ok := seen[key]; ok {
return true
}
}
return false
}
// extendParallelRun returns the exclusive end index of a consecutive parallel-
// safe run starting at start. The run stops before the first call that is not
// parallel-safe or that conflicts on resource keys with any earlier call in
// the same window (so two read_file calls on the same path stay sequential).
func extendParallelRun(registry *tools.Registry, calls []ToolCall, start int, options Options) int {
if start >= len(calls) || !parallelSafeToolCall(registry, calls[start], options) {
return start
}
end := start + 1
keysWindow := [][]string{resourceKeysForCall(registry, calls[start])}
for end < len(calls) {
if !parallelSafeToolCall(registry, calls[end], options) {
break
}
nextKeys := resourceKeysForCall(registry, calls[end])
conflict := false
for _, prev := range keysWindow {
if resourceKeysConflict(prev, nextKeys) {
conflict = true
break
}
}
if conflict {
break
}
keysWindow = append(keysWindow, nextKeys)
end++
}
return end
}
// executeParallelReadBatch runs calls[start:end] concurrently (bounded by
// maxParallelReadTools) and returns results indexed relative to start. All
// execution-side callbacks that can fire inside executeToolCall are serialized
+104
View File
@@ -57,6 +57,20 @@ func (tool *probeTool) Parameters() tools.Schema {
func (tool *probeTool) Safety() tools.Safety {
return tools.Safety{SideEffect: tool.sideEffect, Permission: tools.PermissionAllow, Reason: "test"}
}
// Capabilities maps the probe's SideEffect into the PR5 capability contract so
// parallelSafeToolCall (CapabilitiesOf) can classify reads vs mutators. Without
// this, probes default to EffectUnknown and never enter the concurrent path.
func (tool *probeTool) Capabilities() tools.ToolCapabilities {
switch tool.sideEffect {
case tools.SideEffectRead:
return tools.ToolCapabilities{Effect: tools.EffectReadOnly, ThreadSafe: true}
case tools.SideEffectWrite:
return tools.ToolCapabilities{Effect: tools.EffectWorkspaceWrite, ThreadSafe: false}
default:
return tools.UnknownCapabilities()
}
}
func (tool *probeTool) Run(_ context.Context, args map[string]any) tools.Result {
id, _ := args["id"].(string)
tool.mu.Lock()
@@ -92,6 +106,11 @@ func TestParallelSafeToolCall(t *testing.T) {
registry := tools.NewRegistry()
registry.Register(&probeTool{name: "probe_read", sideEffect: tools.SideEffectRead})
registry.Register(&probeTool{name: "probe_write", sideEffect: tools.SideEffectWrite})
// Read-only but not ThreadSafe: must stay sequential under the PR5 gate.
registry.Register(&keyedProbeTool{
probeTool: probeTool{name: "probe_read_serial", sideEffect: tools.SideEffectRead},
threadSafe: false,
})
call := func(name, args string) ToolCall { return ToolCall{ID: "c", Name: name, Arguments: args} }
if !parallelSafeToolCall(registry, call("probe_read", `{"id":"a"}`), Options{}) {
@@ -109,6 +128,91 @@ func TestParallelSafeToolCall(t *testing.T) {
if parallelSafeToolCall(registry, call("ask_user", `{}`), Options{}) {
t.Fatal("loop-intercepted tools must stay sequential")
}
if parallelSafeToolCall(registry, call("probe_read_serial", `{"id":"a"}`), Options{}) {
t.Fatal("ReadOnly without ThreadSafe must not be parallel-safe")
}
}
func TestResourceKeysConflict(t *testing.T) {
if resourceKeysConflict(nil, []string{"file:a"}) {
t.Fatal("empty vs non-empty must not conflict")
}
if resourceKeysConflict([]string{"file:a"}, nil) {
t.Fatal("non-empty vs empty must not conflict")
}
if resourceKeysConflict([]string{"file:a"}, []string{"file:b"}) {
t.Fatal("distinct keys must not conflict")
}
if !resourceKeysConflict([]string{"file:a", "file:b"}, []string{"file:b"}) {
t.Fatal("shared key must conflict")
}
if resourceKeysConflict([]string{"", "file:a"}, []string{""}) {
t.Fatal("empty-string keys must be ignored")
}
}
func TestExtendParallelRunResourceKeyBoundary(t *testing.T) {
// Two keyed reads on the same path must not share a concurrent window;
// distinct paths may batch together.
registry := tools.NewRegistry()
registry.Register(&keyedProbeTool{
probeTool: probeTool{name: "keyed_read", sideEffect: tools.SideEffectRead},
threadSafe: true,
keys: func(args map[string]any) []string {
id, _ := args["id"].(string)
if id == "" {
return nil
}
return []string{"file:" + id}
},
})
calls := []ToolCall{
{ID: "1", Name: "keyed_read", Arguments: `{"id":"a"}`},
{ID: "2", Name: "keyed_read", Arguments: `{"id":"b"}`},
{ID: "3", Name: "keyed_read", Arguments: `{"id":"a"}`}, // conflicts with call 0
}
// start=0: a and b share no key → run covers [0,2)
if end := extendParallelRun(registry, calls, 0, Options{}); end != 2 {
t.Fatalf("extend from 0 = %d, want 2 (stop before second file:a)", end)
}
// start=2: single remaining call
if end := extendParallelRun(registry, calls, 2, Options{}); end != 3 {
t.Fatalf("extend from 2 = %d, want 3", end)
}
// Empty keys never conflict: three keyless safe reads form one window.
registry.Register(&probeTool{name: "keyless_read", sideEffect: tools.SideEffectRead})
keyless := []ToolCall{
{ID: "1", Name: "keyless_read", Arguments: `{"id":"x"}`},
{ID: "2", Name: "keyless_read", Arguments: `{"id":"y"}`},
{ID: "3", Name: "keyless_read", Arguments: `{"id":"z"}`},
}
if end := extendParallelRun(registry, keyless, 0, Options{}); end != 3 {
t.Fatalf("keyless extend = %d, want 3", end)
}
}
// keyedProbeTool is a probe with explicit ResourceKeys / ThreadSafe for
// planner unit tests.
type keyedProbeTool struct {
probeTool
threadSafe bool
keys func(args map[string]any) []string
}
func (tool *keyedProbeTool) Capabilities() tools.ToolCapabilities {
caps := tools.ToolCapabilities{
Effect: tools.EffectReadOnly,
ThreadSafe: tool.threadSafe,
}
if tool.keys != nil {
caps.ResourceKeys = tool.keys
}
// Honor mutator side effects from the embedded probe when used as write.
if tool.sideEffect == tools.SideEffectWrite {
caps.Effect = tools.EffectWorkspaceWrite
caps.ThreadSafe = false
}
return caps
}
func TestRunExecutesConsecutiveReadsConcurrently(t *testing.T) {
+40 -4
View File
@@ -300,6 +300,18 @@ func TestInteractiveClassifications(t *testing.T) {
func TestReadOnlyClassifications(t *testing.T) {
root := t.TempDir()
// PR6 audited concurrent-safe pure reads (mutex-guarded FileTracker,
// independent scan I/O, independent HTTP). Others stay sequential until
// their shared state is proven safe under concurrent Run.
threadSafeReads := map[string]bool{
"read_file": true,
"read_minified_file": true,
"list_directory": true,
"glob": true,
"grep": true,
"skill": true,
"web_fetch": true,
}
for _, name := range []string{"read_file", "read_minified_file", "list_directory", "glob", "grep", "lsp_navigate", "skill", "web_fetch", "web_search", "tool_search"} {
var found Tool
for _, tool := range BuiltinCatalog(root) {
@@ -315,14 +327,38 @@ func TestReadOnlyClassifications(t *testing.T) {
if caps.Effect != EffectReadOnly {
t.Errorf("%s effect = %v, want ReadOnly", name, caps.Effect)
}
// Current catalog intentionally keeps read tools non-thread-safe
// until each is audited for shared mutable state (FileTracker, etc.).
if caps.ThreadSafe {
t.Errorf("%s ThreadSafe=true without explicit audit", name)
wantSafe := threadSafeReads[name]
if caps.ThreadSafe != wantSafe {
t.Errorf("%s ThreadSafe=%v, want %v", name, caps.ThreadSafe, wantSafe)
}
}
}
func TestScopedScanResourceKeys(t *testing.T) {
// Workspace-wide scan → no keys (do not false-conflict every concurrent scan).
if keys := scopedScanResourceKeys(nil); keys != nil {
t.Fatalf("nil args = %v, want nil", keys)
}
if keys := scopedScanResourceKeys(map[string]any{}); keys != nil {
t.Fatalf("empty args = %v, want nil", keys)
}
if keys := scopedScanResourceKeys(map[string]any{"cwd": "."}); keys != nil {
t.Fatalf("cwd=. = %v, want nil", keys)
}
if keys := scopedScanResourceKeys(map[string]any{"path": "."}); keys != nil {
t.Fatalf("path=. = %v, want nil", keys)
}
// Scoped path → directory: key (glob uses cwd; grep uses path).
keys := scopedScanResourceKeys(map[string]any{"cwd": "internal/tools"})
if len(keys) != 1 || keys[0] != ResourceKeyDirectory+NormalizeResourcePath("internal/tools") {
t.Fatalf("cwd scoped = %v", keys)
}
keys = scopedScanResourceKeys(map[string]any{"path": "./pkg/../pkg/api"})
if len(keys) != 1 || keys[0] != ResourceKeyDirectory+NormalizeResourcePath("./pkg/../pkg/api") {
t.Fatalf("path scoped = %v", keys)
}
}
func TestValidateBuiltinCatalogRejectsUnknownRaw(t *testing.T) {
// Prove the gate is not a rubber stamp: raw Unknown+ThreadSafe is rejected.
problems := ValidateCapabilities("ghost", ToolCapabilities{Effect: EffectUnknown, ThreadSafe: true})
+1 -1
View File
@@ -39,7 +39,7 @@ func NewScopedGlobTool(workspaceRoot string, scope PathScope) Tool {
AdditionalProperties: false,
},
safety: readOnlySafety("Finds matching paths without reading contents or modifying files."),
capabilities: ToolCapabilities{Effect: EffectReadOnly, ThreadSafe: false, ResourceKeys: workspaceResourceKeys},
capabilities: ToolCapabilities{Effect: EffectReadOnly, ThreadSafe: true, ResourceKeys: scopedScanResourceKeys},
},
workspaceRoot: normalizeWorkspaceRoot(workspaceRoot),
scope: scope,
+1 -1
View File
@@ -54,7 +54,7 @@ func NewScopedGrepTool(workspaceRoot string, scope PathScope) Tool {
AdditionalProperties: false,
},
safety: readOnlySafety("Searches file paths and matching lines without modifying files."),
capabilities: ToolCapabilities{Effect: EffectReadOnly, ThreadSafe: false, ResourceKeys: workspaceResourceKeys},
capabilities: ToolCapabilities{Effect: EffectReadOnly, ThreadSafe: true, ResourceKeys: scopedScanResourceKeys},
},
workspaceRoot: normalizeWorkspaceRoot(workspaceRoot),
scope: scope,
+1 -1
View File
@@ -34,7 +34,7 @@ func NewScopedListDirectoryTool(workspaceRoot string, scope PathScope) Tool {
AdditionalProperties: false,
},
safety: readOnlySafety("Lists directory entries without modifying files."),
capabilities: ToolCapabilities{Effect: EffectReadOnly, ThreadSafe: false, ResourceKeys: directoryResourceKeys},
capabilities: ToolCapabilities{Effect: EffectReadOnly, ThreadSafe: true, ResourceKeys: directoryResourceKeys},
},
workspaceRoot: normalizeWorkspaceRoot(workspaceRoot),
scope: scope,
+4 -2
View File
@@ -39,8 +39,10 @@ func NewScopedReadFileTool(workspaceRoot string, scope PathScope) Tool {
AdditionalProperties: false,
},
safety: readOnlySafety("Reads file contents without modifying files."),
// ThreadSafe=false: may update FileTracker session state on read.
capabilities: ToolCapabilities{Effect: EffectReadOnly, ThreadSafe: false, ResourceKeys: fileResourceKeys},
// ThreadSafe: FileTracker is mutex-guarded; concurrent reads of
// distinct paths are safe. Same-path calls still serialize via
// resource-key conflict detection in the agent parallel planner.
capabilities: ToolCapabilities{Effect: EffectReadOnly, ThreadSafe: true, ResourceKeys: fileResourceKeys},
},
workspaceRoot: normalizeWorkspaceRoot(workspaceRoot),
scope: scope,
+1 -1
View File
@@ -34,7 +34,7 @@ func NewScopedReadMinifiedFileTool(workspaceRoot string, scope PathScope) Tool {
AdditionalProperties: false,
},
safety: readOnlySafety("Reads a minified view of file contents without modifying files."),
capabilities: ToolCapabilities{Effect: EffectReadOnly, ThreadSafe: false, ResourceKeys: fileResourceKeys},
capabilities: ToolCapabilities{Effect: EffectReadOnly, ThreadSafe: true, ResourceKeys: fileResourceKeys},
},
workspaceRoot: normalizeWorkspaceRoot(workspaceRoot),
scope: scope,
+14
View File
@@ -127,6 +127,20 @@ func workspaceResourceKeys(_ map[string]any) []string {
return []string{ResourceKeyWorkspace + "root"}
}
// scopedScanResourceKeys returns conflict keys for directory-scoped scan tools
// (glob, grep). Workspace-wide scans (missing / empty / ".") return nil so
// they do not force a false conflict with every other concurrent scan —
// ThreadSafe is the concurrency gate; keys refine same-directory collisions.
// A non-default path/cwd yields a single directory: key.
func scopedScanResourceKeys(args map[string]any) []string {
path := firstStringArg(args, "path", "cwd", "dir", "directory")
normalized := NormalizeResourcePath(path)
if normalized == "" || normalized == "." {
return nil
}
return []string{ResourceKeyDirectory + normalized}
}
// multiFileResourceKeys collects path and paths[] arguments into file: keys.
func multiFileResourceKeys(args map[string]any) []string {
var keys []string
+1 -1
View File
@@ -49,7 +49,7 @@ func NewSkillTool(dir string) *skillTool {
AdditionalProperties: false,
},
safety: readOnlySafety("Reads a local skill file; gathers reusable instructions only."),
capabilities: ToolCapabilities{Effect: EffectReadOnly, ThreadSafe: false},
capabilities: ToolCapabilities{Effect: EffectReadOnly, ThreadSafe: true},
},
}
}
+5 -2
View File
@@ -141,8 +141,11 @@ func newWebFetchToolWithClientAndResolver(client *http.Client, resolver webFetch
Reason: "Fetches remote URL content over the network.",
AdvertiseInAuto: true,
},
// Network read; ThreadSafe=false until client proven concurrent-safe.
capabilities: ToolCapabilities{Effect: EffectReadOnly, ThreadSafe: false, ResourceKeys: endpointResourceKeys},
// Concurrent-safe: net/http.Client is documented as safe for
// concurrent use; each call issues independent request I/O.
// parallelSafeToolCall still requires PermissionAllow (prompted
// network fetches stay sequential until auto-allowed).
capabilities: ToolCapabilities{Effect: EffectReadOnly, ThreadSafe: true, ResourceKeys: endpointResourceKeys},
},
client: client,
resolver: resolver,