feat(extensions): add cancellable interactive tools

commit: c4c30ffcedfc989e3a5af3373e1181e007624317
author: patriceckhart · 38d ago · 2026-09-14 07:42 UTC
parent: f60e492
signature: unsigned
refs: tag: v0.3.73

10 files · +761 -30

M docs/extensions.md +75 -2


@@ -180,6 +180,11 @@ other stdout frame before `hello`.
180	180		"capabilities":["commands","tools","panels"]}
181	181		```
182	182
	183	+	Include `"tool_cancel"` in `capabilities` to opt in to host cancellation
	184	+	notifications. The Go SDK advertises this automatically. The host only sends
	185	+	`tool_cancel` to extensions advertising that exact capability. Missing or
	186	+	unknown capabilities do not enable cancellation notifications.
	187	+
183	188		#### `register_command`
184	189
185	190		```json

@@ -489,6 +494,51 @@ tokens scoped to that extension).
489	494
490	495		`args` is everything the user typed after the command name, trimmed.
491	496
	497	+	#### Interactive tools
	498	+
	499	+	Set `"interactive": true` on `register_tool` for a tool that intentionally
	500	+	waits for user input. Omitted or false retains the normal 60-second reply
	501	+	timeout. Interactive tools have no host-imposed reply deadline, but context
	502	+	cancellation, context deadlines, and extension disconnects still end the call.
	503	+	Print, JSON, RPC, and other hosts without interactive tool support reject
	504	+	these calls with a tool error instead of opening an invisible panel.
	505	+
	506	+	```json
	507	+	{"type":"register_tool","name":"ask_user",
	508	+	"description":"Ask the user a question","schema":{"type":"object"},
	509	+	"interactive":true}
	510	+	```
	511	+
	512	+	The Go SDK provides a context-aware handler:
	513	+
	514	+	```go
	515	+	e.InteractiveTool("ask_user", "Ask the user a question", schema,
	516	+	func(ctx context.Context, args json.RawMessage) ext.ToolResult {
	517	+	// Open your panel and arrange for its callbacks to deliver an answer.
	518	+	// Use invocation-local state, including a unique panel ID.
	519	+	defer closeQuestionPanel()
	520	+	select {
	521	+	case answer := <-answers:
	522	+	return ext.TextResult(answer)
	523	+	case <-ctx.Done():
	524	+	return ext.TextErrorResult("Question cancelled")
	525	+	}
	526	+	})
	527	+	```
	528	+
	529	+	Here `schema`, `answers`, and `closeQuestionPanel` are extension-owned values.
	530	+	Panel dismissal must also resolve the handler's wait. Cancellation is
	531	+	cooperative, the SDK cannot stop a handler that ignores its context. It
	532	+	cancels active handler contexts on `tool_cancel`, shutdown, or host EOF.
	533	+	Existing `Tool` and `DeferredTool` handlers remain source-compatible.
	534	+
	535	+	The registration field is additive and does not change the protocol version.
	536	+	Older hosts ignore it and still apply their normal timeout. Legacy extensions
	537	+	that do not advertise `"tool_cancel"` receive no cancellation frames. The host
	538	+	still stops waiting and discards late results, but cannot request cleanup of
	539	+	abandoned work in those extensions. Both host and extension must support this
	540	+	feature for cancellable, unlimited waits.
	541	+
492	542		#### `tool_call`
493	543
494	544		Sent when the LLM invokes a tool the extension registered. `args` is

@@ -501,8 +551,31 @@ responsible for validating/coercing it.
501	551		```
502	552
503	553		Reply with `tool_result` within the host's tool timeout (default 60s).
504		-	Missing the timeout surfaces an error to the model and the call is
505		-	marked as failed.
	554	+	Interactive tools have no reply timeout. Missing a bounded timeout surfaces
	555	+	an error to the model and marks the call as `timed_out`. Request writes
	556	+	retain a separate bounded transport timeout, even for interactive tools.
	557	+	Cancellation during a blocked write disconnects the extension transport
	558	+	because a partially written frame cannot safely be withdrawn. Other pending
	559	+	calls on that transport also fail.
	560	+
	561	+	#### `tool_cancel`
	562	+
	563	+	Sent best-effort when a delivered tool invocation is abandoned by cancellation
	564	+	or its reply timeout, and only if the extension advertised `"tool_cancel"` in
	565	+	`hello.capabilities`. Normal completion does not send a cancellation frame.
	566	+	`id` is the opaque correlation ID from `tool_call`, not an agent transcript
	567	+	tool ID.
	568	+
	569	+	```json
	570	+	{"type":"tool_cancel","id":"..."}
	571	+	```
	572	+
	573	+	Stop the associated work, close its panel, and discard its pending state.
	574	+	Unknown or already-completed IDs should be ignored. Cancellation can arrive
	575	+	before the handler goroutine starts, so register cancellation state before
	576	+	dispatching work. Delivery is not guaranteed after a transport disconnect.
	577	+	Extensions must also clean up on shutdown and EOF. The host safely ignores
	578	+	late results for abandoned calls, including races with cancellation.
506	579
507	580		#### `event`
508	581


@@ -180,6 +180,11 @@ other stdout frame before `hello`.
180		"capabilities":["commands","tools","panels"]}	180		"capabilities":["commands","tools","panels"]}
181		```	181		```
182			182
			183	+	Include `"tool_cancel"` in `capabilities` to opt in to host cancellation
			184	+	notifications. The Go SDK advertises this automatically. The host only sends
			185	+	`tool_cancel` to extensions advertising that exact capability. Missing or
			186	+	unknown capabilities do not enable cancellation notifications.
			187	+
183		#### `register_command`	188		#### `register_command`
184			189
185		```json	190		```json

@@ -489,6 +494,51 @@ tokens scoped to that extension).
489			494
490		`args` is everything the user typed after the command name, trimmed.	495		`args` is everything the user typed after the command name, trimmed.
491			496
			497	+	#### Interactive tools
			498	+
			499	+	Set `"interactive": true` on `register_tool` for a tool that intentionally
			500	+	waits for user input. Omitted or false retains the normal 60-second reply
			501	+	timeout. Interactive tools have no host-imposed reply deadline, but context
			502	+	cancellation, context deadlines, and extension disconnects still end the call.
			503	+	Print, JSON, RPC, and other hosts without interactive tool support reject
			504	+	these calls with a tool error instead of opening an invisible panel.
			505	+
			506	+	```json
			507	+	{"type":"register_tool","name":"ask_user",
			508	+	"description":"Ask the user a question","schema":{"type":"object"},
			509	+	"interactive":true}
			510	+	```
			511	+
			512	+	The Go SDK provides a context-aware handler:
			513	+
			514	+	```go
			515	+	e.InteractiveTool("ask_user", "Ask the user a question", schema,
			516	+	func(ctx context.Context, args json.RawMessage) ext.ToolResult {
			517	+	// Open your panel and arrange for its callbacks to deliver an answer.
			518	+	// Use invocation-local state, including a unique panel ID.
			519	+	defer closeQuestionPanel()
			520	+	select {
			521	+	case answer := <-answers:
			522	+	return ext.TextResult(answer)
			523	+	case <-ctx.Done():
			524	+	return ext.TextErrorResult("Question cancelled")
			525	+	}
			526	+	})
			527	+	```
			528	+
			529	+	Here `schema`, `answers`, and `closeQuestionPanel` are extension-owned values.
			530	+	Panel dismissal must also resolve the handler's wait. Cancellation is
			531	+	cooperative, the SDK cannot stop a handler that ignores its context. It
			532	+	cancels active handler contexts on `tool_cancel`, shutdown, or host EOF.
			533	+	Existing `Tool` and `DeferredTool` handlers remain source-compatible.
			534	+
			535	+	The registration field is additive and does not change the protocol version.
			536	+	Older hosts ignore it and still apply their normal timeout. Legacy extensions
			537	+	that do not advertise `"tool_cancel"` receive no cancellation frames. The host
			538	+	still stops waiting and discards late results, but cannot request cleanup of
			539	+	abandoned work in those extensions. Both host and extension must support this
			540	+	feature for cancellable, unlimited waits.
			541	+
492		#### `tool_call`	542		#### `tool_call`
493			543
494		Sent when the LLM invokes a tool the extension registered. `args` is	544		Sent when the LLM invokes a tool the extension registered. `args` is

@@ -501,8 +551,31 @@ responsible for validating/coercing it.
501		```	551		```
502			552
503		Reply with `tool_result` within the host's tool timeout (default 60s).	553		Reply with `tool_result` within the host's tool timeout (default 60s).
504	-	Missing the timeout surfaces an error to the model and the call is	554	+	Interactive tools have no reply timeout. Missing a bounded timeout surfaces
505	-	marked as failed.	555	+	an error to the model and marks the call as `timed_out`. Request writes
			556	+	retain a separate bounded transport timeout, even for interactive tools.
			557	+	Cancellation during a blocked write disconnects the extension transport
			558	+	because a partially written frame cannot safely be withdrawn. Other pending
			559	+	calls on that transport also fail.
			560	+
			561	+	#### `tool_cancel`
			562	+
			563	+	Sent best-effort when a delivered tool invocation is abandoned by cancellation
			564	+	or its reply timeout, and only if the extension advertised `"tool_cancel"` in
			565	+	`hello.capabilities`. Normal completion does not send a cancellation frame.
			566	+	`id` is the opaque correlation ID from `tool_call`, not an agent transcript
			567	+	tool ID.
			568	+
			569	+	```json
			570	+	{"type":"tool_cancel","id":"..."}
			571	+	```
			572	+
			573	+	Stop the associated work, close its panel, and discard its pending state.
			574	+	Unknown or already-completed IDs should be ignored. Cancellation can arrive
			575	+	before the handler goroutine starts, so register cancellation state before
			576	+	dispatching work. Delivery is not guaranteed after a transport disconnect.
			577	+	Extensions must also clean up on shutdown and EOF. The host safely ignores
			578	+	late results for abandoned calls, including races with cancellation.
506			579
507		#### `event`	580		#### `event`
508			581

M packages/agent/cli.go +2 -0


@@ -72,6 +72,8 @@ func (h *interactiveExtHooks) ClearNotes(extName string) {
72	72		iv.ClearNotes(extName)
73	73		}
74	74		}
	75	+	func (h *interactiveExtHooks) SupportsInteractiveTools() bool { return h.iv() != nil }
	76	+
75	77		func (h *interactiveExtHooks) OpenPanel(extName string, spec extproto.PanelSpec) {
76	78		if iv := h.iv(); iv != nil {
77	79		iv.OpenPanel(extName, spec)


@@ -72,6 +72,8 @@ func (h *interactiveExtHooks) ClearNotes(extName string) {
72	iv.ClearNotes(extName)	72		iv.ClearNotes(extName)
73	}	73		}
74	}	74		}
		75	+	func (h *interactiveExtHooks) SupportsInteractiveTools() bool { return h.iv() != nil }
		76	+
75	func (h *interactiveExtHooks) OpenPanel(extName string, spec extproto.PanelSpec) {	77		func (h *interactiveExtHooks) OpenPanel(extName string, spec extproto.PanelSpec) {
76	if iv := h.iv(); iv != nil {	78		if iv := h.iv(); iv != nil {
77	iv.OpenPanel(extName, spec)	79		iv.OpenPanel(extName, spec)

M packages/agent/ext/ext.go +57 -8


@@ -33,6 +33,7 @@ package ext
33	33
34	34		import (
35	35		"bufio"
	36	+	"context"
36	37		"encoding/base64"
37	38		"encoding/json"
38	39		"fmt"

@@ -57,6 +58,10 @@ type CommandHandler func(args string) Response
57	58		// ToolResult describing what zot should send back to the model.
58	59		type ToolHandler func(args json.RawMessage) ToolResult
59	60
	61	+	// InteractiveToolHandler waits for input and must release its resources when
	62	+	// ctx is cancelled. Cancellation is cooperative, it cannot stop a goroutine.
	63	+	type InteractiveToolHandler func(ctx context.Context, args json.RawMessage) ToolResult
	64	+
60	65		// EventHandler is called for each lifecycle event the extension
61	66		// subscribed to via Subscribe. The handler is invoked synchronously
62	67		// on the read goroutine; keep it quick or hand off to your own

@@ -300,7 +305,8 @@ type Extension struct {
300	305		mu sync.Mutex
301	306		commands map[string]CommandHandler
302	307		descriptions []descTuple // ordered so register frames arrive in registration order
303		-	tools map[string]ToolHandler
	308	+	tools map[string]InteractiveToolHandler
	309	+	toolCancels map[string]context.CancelFunc
304	310		toolDefs []toolDef // ordered so register frames arrive in registration order
305	311		eventHandlers map[string]EventHandler
306	312		eventNames []string // declared subscription order

@@ -331,6 +337,7 @@ type toolDef struct {
331	337		description string
332	338		schema json.RawMessage
333	339		deferred bool
	340	+	interactive bool
334	341		}
335	342
336	343		// HostInfo is what the host (zot) tells us in HelloAck. Useful for

@@ -355,11 +362,12 @@ func New(name, version string) *Extension {
355	362		out: os.Stdout,
356	363		stderr: os.Stderr,
357	364		commands: map[string]CommandHandler{},
358		-	tools: map[string]ToolHandler{},
	365	+	tools: map[string]InteractiveToolHandler{},
	366	+	toolCancels: map[string]context.CancelFunc{},
359	367		eventHandlers: map[string]EventHandler{},
360	368		panelKeys: map[string]func(key, text string){},
361	369		panelCloses: map[string]func(){},
362		-	caps: []string{"commands", "tools", "events", "panels"},
	370	+	caps: []string{"commands", "tools", "events", "panels", "tool_cancel"},
363	371		}
364	372		}
365	373

@@ -449,9 +457,19 @@ func (e *Extension) DeferredTool(name, description string, schema json.RawMessag
449	457		e.registerTool(name, description, schema, true, fn)
450	458		}
451	459
452		-	func (e *Extension) registerTool(name, description string, schema json.RawMessage, deferred bool, fn ToolHandler) {
	460	+	// InteractiveTool registers a tool that intentionally waits for user input.
	461	+	// The interactive host imposes no reply deadline. Headless hosts reject it.
	462	+	// The handler must honor ctx cancellation and close any panel it owns.
	463	+	func (e *Extension) InteractiveTool(name, description string, schema json.RawMessage, fn InteractiveToolHandler) {
453	464		e.mu.Lock()
	465	+	defer e.mu.Unlock()
454	466		e.tools[name] = fn
	467	+	e.toolDefs = append(e.toolDefs, toolDef{name: name, description: description, schema: schema, interactive: true})
	468	+	}
	469	+
	470	+	func (e *Extension) registerTool(name, description string, schema json.RawMessage, deferred bool, fn ToolHandler) {
	471	+	e.mu.Lock()
	472	+	e.tools[name] = func(_ context.Context, args json.RawMessage) ToolResult { return fn(args) }
455	473		e.toolDefs = append(e.toolDefs, toolDef{name: name, description: description, schema: schema, deferred: deferred})
456	474		e.mu.Unlock()
457	475		}

@@ -538,6 +556,14 @@ func (e *Extension) Notify(level, message string) {
538	556		// Run starts the protocol loop. Blocks until stdin closes (zot has
539	557		// shut us down). Returns the first fatal error, or nil on clean exit.
540	558		func (e *Extension) Run() error {
	559	+	defer func() {
	560	+	e.mu.Lock()
	561	+	defer e.mu.Unlock()
	562	+	for id, cancel := range e.toolCancels {
	563	+	cancel()
	564	+	delete(e.toolCancels, id)
	565	+	}
	566	+	}()
541	567		scanner := bufio.NewScanner(e.in)
542	568		scanner.Buffer(make([]byte, 0, 641024), 41024*1024)
543	569

@@ -604,6 +630,7 @@ func (e *Extension) Run() error {
604	630		Description: td.description,
605	631		Schema: td.schema,
606	632		Deferred: td.deferred,
	633	+	Interactive: td.interactive,
607	634		})
608	635		}
609	636		var intercepts []string

@@ -677,15 +704,37 @@ func (e *Extension) Run() error {
677	704		e.respondTool(tc.ID, TextErrorResult(fmt.Sprintf("no handler for tool %q", tc.Name)))
678	705		continue
679	706		}
680		-	go func(id string, fn ToolHandler, args json.RawMessage) {
	707	+	ctx, cancel := context.WithCancel(context.Background())
	708	+	e.mu.Lock()
	709	+	e.toolCancels[tc.ID] = cancel
	710	+	e.mu.Unlock()
	711	+	go func(id string, fn InteractiveToolHandler, args json.RawMessage) {
681	712		defer func() {
682		-	if r := recover(); r != nil {
	713	+	cancel()
	714	+	e.mu.Lock()
	715	+	delete(e.toolCancels, id)
	716	+	e.mu.Unlock()
	717	+	}()
	718	+	defer func() {
	719	+	if r := recover(); r != nil && ctx.Err() == nil {
683	720		e.respondTool(id, TextErrorResult(fmt.Sprintf("panic: %v", r)))
684	721		}
685	722		}()
686		-	res := fn(args)
687		-	e.respondTool(id, res)
	723	+	res := fn(ctx, args)
	724	+	if ctx.Err() == nil {
	725	+	e.respondTool(id, res)
	726	+	}
688	727		}(tc.ID, fn, tc.Args)
	728	+	case "tool_cancel":
	729	+	var tc extproto.ToolCancelFromHost
	730	+	if json.Unmarshal(line, &tc) != nil {
	731	+	continue
	732	+	}
	733	+	e.mu.Lock()
	734	+	if cancel := e.toolCancels[tc.ID]; cancel != nil {
	735	+	cancel()
	736	+	}
	737	+	e.mu.Unlock()
689	738		case "event":
690	739		var ef extproto.EventFromHost
691	740		if err := json.Unmarshal(line, &ef); err != nil {

 import (
 	"bufio"
+	"context"
 	"encoding/base64"
 	"encoding/json"
 	"fmt"
 // ToolResult describing what zot should send back to the model.
 type ToolHandler func(args json.RawMessage) ToolResult
+// InteractiveToolHandler waits for input and must release its resources when
+// ctx is cancelled. Cancellation is cooperative, it cannot stop a goroutine.
+type InteractiveToolHandler func(ctx context.Context, args json.RawMessage) ToolResult
 // EventHandler is called for each lifecycle event the extension
 // subscribed to via Subscribe. The handler is invoked synchronously
 // on the read goroutine; keep it quick or hand off to your own
 	mu            sync.Mutex
 	commands      map[string]CommandHandler
 	descriptions  []descTuple // ordered so register frames arrive in registration order
-	tools         map[string]ToolHandler
+	tools         map[string]InteractiveToolHandler
+	toolCancels   map[string]context.CancelFunc
 	toolDefs      []toolDef // ordered so register frames arrive in registration order
 	eventHandlers map[string]EventHandler
 	eventNames    []string // declared subscription order
 	description string
 	schema      json.RawMessage
 	deferred    bool
+	interactive bool
 }
 // HostInfo is what the host (zot) tells us in HelloAck. Useful for
 		out:           os.Stdout,
 		stderr:        os.Stderr,
 		commands:      map[string]CommandHandler{},
-		tools:         map[string]ToolHandler{},
+		tools:         map[string]InteractiveToolHandler{},
+		toolCancels:   map[string]context.CancelFunc{},
 		eventHandlers: map[string]EventHandler{},
 		panelKeys:     map[string]func(key, text string){},
 		panelCloses:   map[string]func(){},
-		caps:          []string{"commands", "tools", "events", "panels"},
+		caps:          []string{"commands", "tools", "events", "panels", "tool_cancel"},
 	}
 }
 	e.registerTool(name, description, schema, true, fn)
 }
-func (e *Extension) registerTool(name, description string, schema json.RawMessage, deferred bool, fn ToolHandler) {
+// InteractiveTool registers a tool that intentionally waits for user input.
+// The interactive host imposes no reply deadline. Headless hosts reject it.
+// The handler must honor ctx cancellation and close any panel it owns.
+func (e *Extension) InteractiveTool(name, description string, schema json.RawMessage, fn InteractiveToolHandler) {
 	e.mu.Lock()
+	defer e.mu.Unlock()
 	e.tools[name] = fn
+	e.toolDefs = append(e.toolDefs, toolDef{name: name, description: description, schema: schema, interactive: true})
+}
+func (e *Extension) registerTool(name, description string, schema json.RawMessage, deferred bool, fn ToolHandler) {
+	e.mu.Lock()
+	e.tools[name] = func(_ context.Context, args json.RawMessage) ToolResult { return fn(args) }
 	e.toolDefs = append(e.toolDefs, toolDef{name: name, description: description, schema: schema, deferred: deferred})
 	e.mu.Unlock()
 }
 // Run starts the protocol loop. Blocks until stdin closes (zot has
 // shut us down). Returns the first fatal error, or nil on clean exit.
 func (e *Extension) Run() error {
+	defer func() {
+		e.mu.Lock()
+		defer e.mu.Unlock()
+		for id, cancel := range e.toolCancels {
+			cancel()
+			delete(e.toolCancels, id)
+		}
+	}()
 	scanner := bufio.NewScanner(e.in)
 	scanner.Buffer(make([]byte, 0, 64*1024), 4*1024*1024)
 			Description: td.description,
 			Schema:      td.schema,
 			Deferred:    td.deferred,
+			Interactive: td.interactive,
 		})
 	}
 	var intercepts []string
 				e.respondTool(tc.ID, TextErrorResult(fmt.Sprintf("no handler for tool %q", tc.Name)))
 				continue
 			}
-			go func(id string, fn ToolHandler, args json.RawMessage) {
+			ctx, cancel := context.WithCancel(context.Background())
+			e.mu.Lock()
+			e.toolCancels[tc.ID] = cancel
+			e.mu.Unlock()
+			go func(id string, fn InteractiveToolHandler, args json.RawMessage) {
 				defer func() {
-					if r := recover(); r != nil {
+					cancel()
+					e.mu.Lock()
+					delete(e.toolCancels, id)
+					e.mu.Unlock()
+				}()
+				defer func() {
+					if r := recover(); r != nil && ctx.Err() == nil {
 						e.respondTool(id, TextErrorResult(fmt.Sprintf("panic: %v", r)))
 					}
 				}()
-				res := fn(args)
+				res := fn(ctx, args)
-				e.respondTool(id, res)
+				if ctx.Err() == nil {
+					e.respondTool(id, res)
+				}
 			}(tc.ID, fn, tc.Args)
+		case "tool_cancel":
+			var tc extproto.ToolCancelFromHost
+			if json.Unmarshal(line, &tc) != nil {
+				continue
+			}
+			e.mu.Lock()
+			if cancel := e.toolCancels[tc.ID]; cancel != nil {
+				cancel()
+			}
+			e.mu.Unlock()
 		case "event":
 			var ef extproto.EventFromHost
 			if err := json.Unmarshal(line, &ef); err != nil {

A packages/agent/ext/interactive_tool_test.go +115 -0

@@ -0,0 +1,115 @@
	1	+	package ext
	2	+
	3	+	import (
	4	+	"context"
	5	+	"encoding/json"
	6	+	"io"
	7	+	"slices"
	8	+	"testing"
	9	+	"time"
	10	+
	11	+	"github.com/patriceckhart/zot/packages/agent/extproto"
	12	+	)
	13	+
	14	+	func TestInteractiveToolLifecycle(t *testing.T) {
	15	+	for _, end := range []string{"cancel", "shutdown", "eof", "complete"} {
	16	+	t.Run(end, func(t *testing.T) {
	17	+	h := newHarness("interactive")
	18	+	started := make(chan context.Context, 2)
	19	+	finished := make(chan struct{}, 2)
	20	+	answer := make(chan struct{})
	21	+	h.ext.InteractiveTool("ask", "Ask the user", json.RawMessage(`{}`), func(ctx context.Context, _ json.RawMessage) ToolResult {
	22	+	started <- ctx
	23	+	defer func() { finished <- struct{}{} }()
	24	+	select {
	25	+	case <-ctx.Done():
	26	+	return TextErrorResult("cancelled")
	27	+	case <-answer:
	28	+	return TextResult("answer")
	29	+	}
	30	+	})
	31	+	done := make(chan error, 1)
	32	+	go func() { done <- h.ext.Run() }()
	33	+	t.Cleanup(func() {
	34	+	h.hostW.Close()
	35	+	select {
	36	+	case err := <-done:
	37	+	if err != nil {
	38	+	t.Error(err)
	39	+	}
	40	+	case <-time.After(time.Second):
	41	+	t.Error("SDK did not exit")
	42	+	}
	43	+	h.ext.out.(io.Closer).Close()
	44	+	h.ext.in.(io.Closer).Close()
	45	+	})
	46	+	var hello extproto.HelloFromExt
	47	+	if err := json.Unmarshal(h.next(t).raw, &hello); err != nil {
	48	+	t.Fatal(err)
	49	+	}
	50	+	if hello.Type != "hello" \|\| !slices.Contains(hello.Capabilities, "tool_cancel") {
	51	+	t.Fatalf("SDK did not advertise cancellation support: %+v", hello)
	52	+	}
	53	+	h.sendToExt(t, extproto.HelloAckFromHost{Type: "hello_ack"})
	54	+	var registration extproto.RegisterToolFromExt
	55	+	json.Unmarshal(h.drainUntil(t, "register_tool").raw, &registration)
	56	+	if !registration.Interactive \|\| registration.Name != "ask" {
	57	+	t.Fatalf("registration: %+v", registration)
	58	+	}
	59	+	h.drainUntil(t, "ready")
	60	+	h.sendToExt(t, extproto.ToolCallFromHost{Type: "tool_call", ID: "one", Name: "ask", Args: json.RawMessage(`{}`)})
	61	+	// Send cancellation without waiting for the goroutine to start.
	62	+	if end == "cancel" {
	63	+	h.sendToExt(t, extproto.ToolCancelFromHost{Type: "tool_cancel", ID: "one"})
	64	+	}
	65	+	var ctx context.Context
	66	+	select {
	67	+	case ctx = <-started:
	68	+	case <-time.After(time.Second):
	69	+	t.Fatal("handler did not start")
	70	+	}
	71	+	switch end {
	72	+	case "shutdown":
	73	+	h.sendToExt(t, extproto.ShutdownFromHost{Type: "shutdown"})
	74	+	h.drainUntil(t, "shutdown_ack")
	75	+	case "eof":
	76	+	h.hostW.Close()
	77	+	case "complete":
	78	+	close(answer)
	79	+	var result extproto.ToolResultFromExt
	80	+	json.Unmarshal(h.drainUntil(t, "tool_result").raw, &result)
	81	+	if result.ID != "one" \|\| len(result.Content) != 1 \|\| result.Content[0].Text != "answer" {
	82	+	t.Fatalf("result: %+v", result)
	83	+	}
	84	+	}
	85	+	select {
	86	+	case <-finished:
	87	+	case <-time.After(time.Second):
	88	+	t.Fatal("handler did not finish")
	89	+	}
	90	+	if end != "complete" && ctx.Err() != context.Canceled {
	91	+	t.Fatal("handler context was not cancelled")
	92	+	}
	93	+	if end == "cancel" {
	94	+	// Unknown and duplicate cancellation must not affect another call.
	95	+	h.sendToExt(t, extproto.ToolCancelFromHost{Type: "tool_cancel", ID: "one"})
	96	+	h.sendToExt(t, extproto.ToolCancelFromHost{Type: "tool_cancel", ID: "unknown"})
	97	+	h.sendToExt(t, extproto.ToolCallFromHost{Type: "tool_call", ID: "two", Name: "ask", Args: json.RawMessage(`{}`)})
	98	+	select {
	99	+	case ctx = <-started:
	100	+	case <-time.After(time.Second):
	101	+	t.Fatal("second handler did not start")
	102	+	}
	103	+	if ctx.Err() != nil {
	104	+	t.Fatal("second call was cancelled")
	105	+	}
	106	+	close(answer)
	107	+	var result extproto.ToolResultFromExt
	108	+	json.Unmarshal(h.drainUntil(t, "tool_result").raw, &result)
	109	+	if result.ID != "two" {
	110	+	t.Fatalf("unexpected result: %+v", result)
	111	+	}
	112	+	}
	113	+	})
	114	+	}
	115	+	}

@@ -0,0 +1,115 @@
	1	+	package ext
	2	+
	3	+	import (
	4	+	"context"
	5	+	"encoding/json"
	6	+	"io"
	7	+	"slices"
	8	+	"testing"
	9	+	"time"
	10	+
	11	+	"github.com/patriceckhart/zot/packages/agent/extproto"
	12	+	)
	13	+
	14	+	func TestInteractiveToolLifecycle(t *testing.T) {
	15	+	for _, end := range []string{"cancel", "shutdown", "eof", "complete"} {
	16	+	t.Run(end, func(t *testing.T) {
	17	+	h := newHarness("interactive")
	18	+	started := make(chan context.Context, 2)
	19	+	finished := make(chan struct{}, 2)
	20	+	answer := make(chan struct{})
	21	+	h.ext.InteractiveTool("ask", "Ask the user", json.RawMessage(`{}`), func(ctx context.Context, _ json.RawMessage) ToolResult {
	22	+	started <- ctx
	23	+	defer func() { finished <- struct{}{} }()
	24	+	select {
	25	+	case <-ctx.Done():
	26	+	return TextErrorResult("cancelled")
	27	+	case <-answer:
	28	+	return TextResult("answer")
	29	+	}
	30	+	})
	31	+	done := make(chan error, 1)
	32	+	go func() { done <- h.ext.Run() }()
	33	+	t.Cleanup(func() {
	34	+	h.hostW.Close()
	35	+	select {
	36	+	case err := <-done:
	37	+	if err != nil {
	38	+	t.Error(err)
	39	+	}
	40	+	case <-time.After(time.Second):
	41	+	t.Error("SDK did not exit")
	42	+	}
	43	+	h.ext.out.(io.Closer).Close()
	44	+	h.ext.in.(io.Closer).Close()
	45	+	})
	46	+	var hello extproto.HelloFromExt
	47	+	if err := json.Unmarshal(h.next(t).raw, &hello); err != nil {
	48	+	t.Fatal(err)
	49	+	}
	50	+	if hello.Type != "hello" \|\| !slices.Contains(hello.Capabilities, "tool_cancel") {
	51	+	t.Fatalf("SDK did not advertise cancellation support: %+v", hello)
	52	+	}
	53	+	h.sendToExt(t, extproto.HelloAckFromHost{Type: "hello_ack"})
	54	+	var registration extproto.RegisterToolFromExt
	55	+	json.Unmarshal(h.drainUntil(t, "register_tool").raw, &registration)
	56	+	if !registration.Interactive \|\| registration.Name != "ask" {
	57	+	t.Fatalf("registration: %+v", registration)
	58	+	}
	59	+	h.drainUntil(t, "ready")
	60	+	h.sendToExt(t, extproto.ToolCallFromHost{Type: "tool_call", ID: "one", Name: "ask", Args: json.RawMessage(`{}`)})
	61	+	// Send cancellation without waiting for the goroutine to start.
	62	+	if end == "cancel" {
	63	+	h.sendToExt(t, extproto.ToolCancelFromHost{Type: "tool_cancel", ID: "one"})
	64	+	}
	65	+	var ctx context.Context
	66	+	select {
	67	+	case ctx = <-started:
	68	+	case <-time.After(time.Second):
	69	+	t.Fatal("handler did not start")
	70	+	}
	71	+	switch end {
	72	+	case "shutdown":
	73	+	h.sendToExt(t, extproto.ShutdownFromHost{Type: "shutdown"})
	74	+	h.drainUntil(t, "shutdown_ack")
	75	+	case "eof":
	76	+	h.hostW.Close()
	77	+	case "complete":
	78	+	close(answer)
	79	+	var result extproto.ToolResultFromExt
	80	+	json.Unmarshal(h.drainUntil(t, "tool_result").raw, &result)
	81	+	if result.ID != "one" \|\| len(result.Content) != 1 \|\| result.Content[0].Text != "answer" {
	82	+	t.Fatalf("result: %+v", result)
	83	+	}
	84	+	}
	85	+	select {
	86	+	case <-finished:
	87	+	case <-time.After(time.Second):
	88	+	t.Fatal("handler did not finish")
	89	+	}
	90	+	if end != "complete" && ctx.Err() != context.Canceled {
	91	+	t.Fatal("handler context was not cancelled")
	92	+	}
	93	+	if end == "cancel" {
	94	+	// Unknown and duplicate cancellation must not affect another call.
	95	+	h.sendToExt(t, extproto.ToolCancelFromHost{Type: "tool_cancel", ID: "one"})
	96	+	h.sendToExt(t, extproto.ToolCancelFromHost{Type: "tool_cancel", ID: "unknown"})
	97	+	h.sendToExt(t, extproto.ToolCallFromHost{Type: "tool_call", ID: "two", Name: "ask", Args: json.RawMessage(`{}`)})
	98	+	select {
	99	+	case ctx = <-started:
	100	+	case <-time.After(time.Second):
	101	+	t.Fatal("second handler did not start")
	102	+	}
	103	+	if ctx.Err() != nil {
	104	+	t.Fatal("second call was cancelled")
	105	+	}
	106	+	close(answer)
	107	+	var result extproto.ToolResultFromExt
	108	+	json.Unmarshal(h.drainUntil(t, "tool_result").raw, &result)
	109	+	if result.ID != "two" {
	110	+	t.Fatalf("unexpected result: %+v", result)
	111	+	}
	112	+	}
	113	+	})
	114	+	}
	115	+	}

A packages/agent/extensions/interactive_tool_test.go +287 -0

@@ -0,0 +1,287 @@
	1	+	package extensions
	2	+
	3	+	import (
	4	+	"bufio"
	5	+	"context"
	6	+	"encoding/json"
	7	+	"errors"
	8	+	"io"
	9	+	"os"
	10	+	"sync"
	11	+	"testing"
	12	+	"testing/synctest"
	13	+	"time"
	14	+
	15	+	"github.com/patriceckhart/zot/packages/agent/extproto"
	16	+	)
	17	+
	18	+	type interactiveHooks struct{ stubHooks }
	19	+
	20	+	func (*interactiveHooks) SupportsInteractiveTools() bool { return true }
	21	+
	22	+	func toolPeer(t testing.T) (Manager, Extension, <-chan []byte, io.PipeWriter) {
	23	+	t.Helper()
	24	+	m := New(t.TempDir(), "", "test", "", "", &interactiveHooks{})
	25	+	frames, ext := lifecyclePeer(t, m)
	26	+	ext.pendingTool = map[string]chan extproto.ToolResultFromExt{}
	27	+	ext.toolCancel = true
	28	+	ext.readyCh = make(chan struct{})
	29	+	log, err := os.CreateTemp(t.TempDir(), "log")
	30	+	if err != nil {
	31	+	t.Fatal(err)
	32	+	}
	33	+	ext.logFile = log
	34	+	r, w := io.Pipe()
	35	+	done := make(chan struct{})
	36	+	go func() {
	37	+	m.readLoop(ext, bufio.NewScanner(r))
	38	+	close(done)
	39	+	}()
	40	+	t.Cleanup(func() {
	41	+	w.Close()
	42	+	<-done
	43	+	r.Close()
	44	+	log.Close()
	45	+	})
	46	+	m.toolIndex["ask"] = ext
	47	+	return m, ext, frames, w
	48	+	}
	49	+
	50	+	func sendToolFrame(t *testing.T, w io.Writer, frame any) {
	51	+	t.Helper()
	52	+	b, err := extproto.Encode(frame)
	53	+	if err != nil {
	54	+	t.Fatal(err)
	55	+	}
	56	+	if _, err := w.Write(b); err != nil {
	57	+	t.Fatal(err)
	58	+	}
	59	+	}
	60	+
	61	+	func toolCallFrame(t *testing.T, frames <-chan []byte) extproto.ToolCallFromHost {
	62	+	t.Helper()
	63	+	var call extproto.ToolCallFromHost
	64	+	if err := json.Unmarshal(nextLifecycleFrame(t, frames), &call); err != nil {
	65	+	t.Fatal(err)
	66	+	}
	67	+	if call.Type != "tool_call" {
	68	+	t.Fatalf("expected tool_call, got %+v", call)
	69	+	}
	70	+	return call
	71	+	}
	72	+
	73	+	func awaitToolError(t *testing.T, done <-chan error) error {
	74	+	t.Helper()
	75	+	select {
	76	+	case err := <-done:
	77	+	return err
	78	+	case <-time.After(2 * time.Second):
	79	+	t.Fatal("invocation did not finish")
	80	+	return nil
	81	+	}
	82	+	}
	83	+
	84	+	func TestInteractiveToolRegistrationAndCompletion(t *testing.T) {
	85	+	m, _, frames, replies := toolPeer(t)
	86	+	sendToolFrame(t, replies, extproto.RegisterToolFromExt{Type: "register_tool", Name: "ask", Schema: json.RawMessage(`{}`), Interactive: true})
	87	+	sendToolFrame(t, replies, extproto.ReadyFromExt{Type: "ready"})
	88	+	m.WaitForReady(time.Second)
	89	+	infos := m.Tools()
	90	+	if len(infos) != 1 \|\| !infos[0].Interactive {
	91	+	t.Fatalf("registration: %+v", infos)
	92	+	}
	93	+	tool := NewTool(m, infos[0]).(*extensionTool)
	94	+	if tool.timeout != 0 {
	95	+	t.Fatal("interactive timeout is enabled")
	96	+	}
	97	+	if normal := NewTool(m, ToolInfo{}).(extensionTool); normal.timeout != 60time.Second {
	98	+	t.Fatal("normal timeout changed")
	99	+	}
	100	+	done := make(chan error, 1)
	101	+	go func() {
	102	+	result, err := tool.Execute(context.Background(), nil, nil)
	103	+	if err == nil && result.IsError {
	104	+	err = errors.New("interactive completion failed")
	105	+	}
	106	+	done <- err
	107	+	}()
	108	+	call := toolCallFrame(t, frames)
	109	+	sendToolFrame(t, replies, extproto.ToolResultFromExt{Type: "tool_result", ID: call.ID})
	110	+	if err := awaitToolError(t, done); err != nil {
	111	+	t.Fatal(err)
	112	+	}
	113	+	}
	114	+
	115	+	func TestInteractiveToolTermination(t *testing.T) {
	116	+	for _, reason := range []string{"cancel", "timeout", "disconnect", "shutdown"} {
	117	+	t.Run(reason, func(t *testing.T) {
	118	+	synctest.Test(t, func(t *testing.T) {
	119	+	m, ext, frames, replies := toolPeer(t)
	120	+	ctx, cancel := context.WithCancel(context.Background())
	121	+	defer cancel()
	122	+	timeout := time.Duration(0)
	123	+	if reason == "timeout" {
	124	+	timeout = time.Nanosecond
	125	+	}
	126	+	done := make(chan error, 1)
	127	+	go func() {
	128	+	_, err := m.InvokeTool(ctx, "ask", nil, timeout)
	129	+	done <- err
	130	+	}()
	131	+	call := toolCallFrame(t, frames)
	132	+	// Ensure request delivery has completed before cancelling the reply wait.
	133	+	synctest.Wait()
	134	+	switch reason {
	135	+	case "cancel":
	136	+	cancel()
	137	+	case "disconnect":
	138	+	replies.Close()
	139	+	case "shutdown":
	140	+	m.Stop(time.Second)
	141	+	}
	142	+	err := awaitToolError(t, done)
	143	+	if err == nil {
	144	+	t.Fatal("expected terminal error")
	145	+	}
	146	+	if reason == "cancel" && !errors.Is(err, context.Canceled) {
	147	+	t.Fatal(err)
	148	+	}
	149	+	if reason == "timeout" && !errors.Is(err, context.DeadlineExceeded) {
	150	+	t.Fatal(err)
	151	+	}
	152	+	ext.mu.Lock()
	153	+	pending := len(ext.pendingTool)
	154	+	ext.mu.Unlock()
	155	+	if pending != 0 {
	156	+	t.Fatal("pending call leaked")
	157	+	}
	158	+	if reason == "cancel" \|\| reason == "timeout" {
	159	+	var frame extproto.ToolCancelFromHost
	160	+	json.Unmarshal(nextLifecycleFrame(t, frames), &frame)
	161	+	if frame.Type != "tool_cancel" \|\| frame.ID != call.ID {
	162	+	t.Fatalf("cancel: %+v", frame)
	163	+	}
	164	+	// A late result must not block delivery of the next result.
	165	+	sendToolFrame(t, replies, extproto.ToolResultFromExt{Type: "tool_result", ID: call.ID})
	166	+	go func() {
	167	+	_, err := m.InvokeTool(context.Background(), "ask", nil, 0)
	168	+	done <- err
	169	+	}()
	170	+	next := toolCallFrame(t, frames)
	171	+	sendToolFrame(t, replies, extproto.ToolResultFromExt{Type: "tool_result", ID: next.ID})
	172	+	if err := awaitToolError(t, done); err != nil {
	173	+	t.Fatal(err)
	174	+	}
	175	+	}
	176	+	})
	177	+	})
	178	+	}
	179	+	}
	180	+
	181	+	func TestInteractiveToolConcurrentReplies(t *testing.T) {
	182	+	m, _, frames, replies := toolPeer(t)
	183	+	const count = 32
	184	+	done := make(chan error, count)
	185	+	for n := 0; n < count; n++ {
	186	+	go func() {
	187	+	_, err := m.InvokeTool(context.Background(), "ask", nil, 0)
	188	+	done <- err
	189	+	}()
	190	+	}
	191	+	ids := make([]string, count)
	192	+	seen := map[string]bool{}
	193	+	for n := range ids {
	194	+	ids[n] = toolCallFrame(t, frames).ID
	195	+	if seen[ids[n]] {
	196	+	t.Fatal("duplicate invocation ID")
	197	+	}
	198	+	seen[ids[n]] = true
	199	+	}
	200	+	for n := count - 1; n >= 0; n-- {
	201	+	sendToolFrame(t, replies, extproto.ToolResultFromExt{Type: "tool_result", ID: ids[n]})
	202	+	}
	203	+	for range ids {
	204	+	if err := awaitToolError(t, done); err != nil {
	205	+	t.Fatal(err)
	206	+	}
	207	+	}
	208	+	}
	209	+
	210	+	func TestInteractiveToolWaitsBeyondDefaultTimeout(t *testing.T) {
	211	+	synctest.Test(t, func(t *testing.T) {
	212	+	m, ext, frames, replies := toolPeer(t)
	213	+	done := make(chan error, 1)
	214	+	go func() {
	215	+	_, err := m.InvokeTool(context.Background(), "ask", nil, 0)
	216	+	done <- err
	217	+	}()
	218	+	call := toolCallFrame(t, frames)
	219	+	synctest.Wait()
	220	+	// Advance the fake clock, not wall-clock time.
	221	+	time.Sleep(2 * time.Minute)
	222	+	synctest.Wait()
	223	+	select {
	224	+	case err := <-done:
	225	+	t.Fatalf("interactive call ended early: %v", err)
	226	+	default:
	227	+	}
	228	+	ext.mu.Lock()
	229	+	pending := len(ext.pendingTool)
	230	+	ext.mu.Unlock()
	231	+	if pending != 1 {
	232	+	t.Fatal("interactive call is no longer pending")
	233	+	}
	234	+	sendToolFrame(t, replies, extproto.ToolResultFromExt{Type: "tool_result", ID: call.ID})
	235	+	if err := awaitToolError(t, done); err != nil {
	236	+	t.Fatal(err)
	237	+	}
	238	+	})
	239	+	}
	240	+
	241	+	func TestInteractiveToolHeadless(t *testing.T) {
	242	+	m := New(t.TempDir(), "", "test", "", "", &stubHooks{})
	243	+	tool := NewTool(m, ToolInfo{Name: "ask", Interactive: true})
	244	+	result, err := tool.Execute(context.Background(), nil, nil)
	245	+	if err != nil \|\| !result.IsError \|\| result.Status != "failed" {
	246	+	t.Fatalf("headless result: %+v, %v", result, err)
	247	+	}
	248	+	}
	249	+
	250	+	type blockedToolWriter struct {
	251	+	started chan struct{}
	252	+	closed chan struct{}
	253	+	once sync.Once
	254	+	}
	255	+
	256	+	func (w *blockedToolWriter) Write([]byte) (int, error) {
	257	+	close(w.started)
	258	+	<-w.closed
	259	+	return 0, io.ErrClosedPipe
	260	+	}
	261	+
	262	+	func (w *blockedToolWriter) Close() error {
	263	+	w.once.Do(func() { close(w.closed) })
	264	+	return nil
	265	+	}
	266	+
	267	+	func TestToolCancellationDuringWrite(t *testing.T) {
	268	+	w := &blockedToolWriter{started: make(chan struct{}), closed: make(chan struct{})}
	269	+	pipe := newOrderedPipe(w)
	270	+	defer pipe.Close()
	271	+	ctx, cancel := context.WithCancel(context.Background())
	272	+	defer cancel()
	273	+	done := make(chan error, 1)
	274	+	go func() {
	275	+	_, err := pipe.writeContext(ctx, []byte("{}\n"), time.Hour)
	276	+	done <- err
	277	+	}()
	278	+	select {
	279	+	case <-w.started:
	280	+	case <-time.After(time.Second):
	281	+	t.Fatal("write did not start")
	282	+	}
	283	+	cancel()
	284	+	if err := awaitToolError(t, done); !errors.Is(err, context.Canceled) {
	285	+	t.Fatalf("write cancellation: %v", err)
	286	+	}
	287	+	}

@@ -0,0 +1,287 @@
	1	+	package extensions
	2	+
	3	+	import (
	4	+	"bufio"
	5	+	"context"
	6	+	"encoding/json"
	7	+	"errors"
	8	+	"io"
	9	+	"os"
	10	+	"sync"
	11	+	"testing"
	12	+	"testing/synctest"
	13	+	"time"
	14	+
	15	+	"github.com/patriceckhart/zot/packages/agent/extproto"
	16	+	)
	17	+
	18	+	type interactiveHooks struct{ stubHooks }
	19	+
	20	+	func (*interactiveHooks) SupportsInteractiveTools() bool { return true }
	21	+
	22	+	func toolPeer(t testing.T) (Manager, Extension, <-chan []byte, io.PipeWriter) {
	23	+	t.Helper()
	24	+	m := New(t.TempDir(), "", "test", "", "", &interactiveHooks{})
	25	+	frames, ext := lifecyclePeer(t, m)
	26	+	ext.pendingTool = map[string]chan extproto.ToolResultFromExt{}
	27	+	ext.toolCancel = true
	28	+	ext.readyCh = make(chan struct{})
	29	+	log, err := os.CreateTemp(t.TempDir(), "log")
	30	+	if err != nil {
	31	+	t.Fatal(err)
	32	+	}
	33	+	ext.logFile = log
	34	+	r, w := io.Pipe()
	35	+	done := make(chan struct{})
	36	+	go func() {
	37	+	m.readLoop(ext, bufio.NewScanner(r))
	38	+	close(done)
	39	+	}()
	40	+	t.Cleanup(func() {
	41	+	w.Close()
	42	+	<-done
	43	+	r.Close()
	44	+	log.Close()
	45	+	})
	46	+	m.toolIndex["ask"] = ext
	47	+	return m, ext, frames, w
	48	+	}
	49	+
	50	+	func sendToolFrame(t *testing.T, w io.Writer, frame any) {
	51	+	t.Helper()
	52	+	b, err := extproto.Encode(frame)
	53	+	if err != nil {
	54	+	t.Fatal(err)
	55	+	}
	56	+	if _, err := w.Write(b); err != nil {
	57	+	t.Fatal(err)
	58	+	}
	59	+	}
	60	+
	61	+	func toolCallFrame(t *testing.T, frames <-chan []byte) extproto.ToolCallFromHost {
	62	+	t.Helper()
	63	+	var call extproto.ToolCallFromHost
	64	+	if err := json.Unmarshal(nextLifecycleFrame(t, frames), &call); err != nil {
	65	+	t.Fatal(err)
	66	+	}
	67	+	if call.Type != "tool_call" {
	68	+	t.Fatalf("expected tool_call, got %+v", call)
	69	+	}
	70	+	return call
	71	+	}
	72	+
	73	+	func awaitToolError(t *testing.T, done <-chan error) error {
	74	+	t.Helper()
	75	+	select {
	76	+	case err := <-done:
	77	+	return err
	78	+	case <-time.After(2 * time.Second):
	79	+	t.Fatal("invocation did not finish")
	80	+	return nil
	81	+	}
	82	+	}
	83	+
	84	+	func TestInteractiveToolRegistrationAndCompletion(t *testing.T) {
	85	+	m, _, frames, replies := toolPeer(t)
	86	+	sendToolFrame(t, replies, extproto.RegisterToolFromExt{Type: "register_tool", Name: "ask", Schema: json.RawMessage(`{}`), Interactive: true})
	87	+	sendToolFrame(t, replies, extproto.ReadyFromExt{Type: "ready"})
	88	+	m.WaitForReady(time.Second)
	89	+	infos := m.Tools()
	90	+	if len(infos) != 1 \|\| !infos[0].Interactive {
	91	+	t.Fatalf("registration: %+v", infos)
	92	+	}
	93	+	tool := NewTool(m, infos[0]).(*extensionTool)
	94	+	if tool.timeout != 0 {
	95	+	t.Fatal("interactive timeout is enabled")
	96	+	}
	97	+	if normal := NewTool(m, ToolInfo{}).(extensionTool); normal.timeout != 60time.Second {
	98	+	t.Fatal("normal timeout changed")
	99	+	}
	100	+	done := make(chan error, 1)
	101	+	go func() {
	102	+	result, err := tool.Execute(context.Background(), nil, nil)
	103	+	if err == nil && result.IsError {
	104	+	err = errors.New("interactive completion failed")
	105	+	}
	106	+	done <- err
	107	+	}()
	108	+	call := toolCallFrame(t, frames)
	109	+	sendToolFrame(t, replies, extproto.ToolResultFromExt{Type: "tool_result", ID: call.ID})
	110	+	if err := awaitToolError(t, done); err != nil {
	111	+	t.Fatal(err)
	112	+	}
	113	+	}
	114	+
	115	+	func TestInteractiveToolTermination(t *testing.T) {
	116	+	for _, reason := range []string{"cancel", "timeout", "disconnect", "shutdown"} {
	117	+	t.Run(reason, func(t *testing.T) {
	118	+	synctest.Test(t, func(t *testing.T) {
	119	+	m, ext, frames, replies := toolPeer(t)
	120	+	ctx, cancel := context.WithCancel(context.Background())
	121	+	defer cancel()
	122	+	timeout := time.Duration(0)
	123	+	if reason == "timeout" {
	124	+	timeout = time.Nanosecond
	125	+	}
	126	+	done := make(chan error, 1)
	127	+	go func() {
	128	+	_, err := m.InvokeTool(ctx, "ask", nil, timeout)
	129	+	done <- err
	130	+	}()
	131	+	call := toolCallFrame(t, frames)
	132	+	// Ensure request delivery has completed before cancelling the reply wait.
	133	+	synctest.Wait()
	134	+	switch reason {
	135	+	case "cancel":
	136	+	cancel()
	137	+	case "disconnect":
	138	+	replies.Close()
	139	+	case "shutdown":
	140	+	m.Stop(time.Second)
	141	+	}
	142	+	err := awaitToolError(t, done)
	143	+	if err == nil {
	144	+	t.Fatal("expected terminal error")
	145	+	}
	146	+	if reason == "cancel" && !errors.Is(err, context.Canceled) {
	147	+	t.Fatal(err)
	148	+	}
	149	+	if reason == "timeout" && !errors.Is(err, context.DeadlineExceeded) {
	150	+	t.Fatal(err)
	151	+	}
	152	+	ext.mu.Lock()
	153	+	pending := len(ext.pendingTool)
	154	+	ext.mu.Unlock()
	155	+	if pending != 0 {
	156	+	t.Fatal("pending call leaked")
	157	+	}
	158	+	if reason == "cancel" \|\| reason == "timeout" {
	159	+	var frame extproto.ToolCancelFromHost
	160	+	json.Unmarshal(nextLifecycleFrame(t, frames), &frame)
	161	+	if frame.Type != "tool_cancel" \|\| frame.ID != call.ID {
	162	+	t.Fatalf("cancel: %+v", frame)
	163	+	}
	164	+	// A late result must not block delivery of the next result.
	165	+	sendToolFrame(t, replies, extproto.ToolResultFromExt{Type: "tool_result", ID: call.ID})
	166	+	go func() {
	167	+	_, err := m.InvokeTool(context.Background(), "ask", nil, 0)
	168	+	done <- err
	169	+	}()
	170	+	next := toolCallFrame(t, frames)
	171	+	sendToolFrame(t, replies, extproto.ToolResultFromExt{Type: "tool_result", ID: next.ID})
	172	+	if err := awaitToolError(t, done); err != nil {
	173	+	t.Fatal(err)
	174	+	}
	175	+	}
	176	+	})
	177	+	})
	178	+	}
	179	+	}
	180	+
	181	+	func TestInteractiveToolConcurrentReplies(t *testing.T) {
	182	+	m, _, frames, replies := toolPeer(t)
	183	+	const count = 32
	184	+	done := make(chan error, count)
	185	+	for n := 0; n < count; n++ {
	186	+	go func() {
	187	+	_, err := m.InvokeTool(context.Background(), "ask", nil, 0)
	188	+	done <- err
	189	+	}()
	190	+	}
	191	+	ids := make([]string, count)
	192	+	seen := map[string]bool{}
	193	+	for n := range ids {
	194	+	ids[n] = toolCallFrame(t, frames).ID
	195	+	if seen[ids[n]] {
	196	+	t.Fatal("duplicate invocation ID")
	197	+	}
	198	+	seen[ids[n]] = true
	199	+	}
	200	+	for n := count - 1; n >= 0; n-- {
	201	+	sendToolFrame(t, replies, extproto.ToolResultFromExt{Type: "tool_result", ID: ids[n]})
	202	+	}
	203	+	for range ids {
	204	+	if err := awaitToolError(t, done); err != nil {
	205	+	t.Fatal(err)
	206	+	}
	207	+	}
	208	+	}
	209	+
	210	+	func TestInteractiveToolWaitsBeyondDefaultTimeout(t *testing.T) {
	211	+	synctest.Test(t, func(t *testing.T) {
	212	+	m, ext, frames, replies := toolPeer(t)
	213	+	done := make(chan error, 1)
	214	+	go func() {
	215	+	_, err := m.InvokeTool(context.Background(), "ask", nil, 0)
	216	+	done <- err
	217	+	}()
	218	+	call := toolCallFrame(t, frames)
	219	+	synctest.Wait()
	220	+	// Advance the fake clock, not wall-clock time.
	221	+	time.Sleep(2 * time.Minute)
	222	+	synctest.Wait()
	223	+	select {
	224	+	case err := <-done:
	225	+	t.Fatalf("interactive call ended early: %v", err)
	226	+	default:
	227	+	}
	228	+	ext.mu.Lock()
	229	+	pending := len(ext.pendingTool)
	230	+	ext.mu.Unlock()
	231	+	if pending != 1 {
	232	+	t.Fatal("interactive call is no longer pending")
	233	+	}
	234	+	sendToolFrame(t, replies, extproto.ToolResultFromExt{Type: "tool_result", ID: call.ID})
	235	+	if err := awaitToolError(t, done); err != nil {
	236	+	t.Fatal(err)
	237	+	}
	238	+	})
	239	+	}
	240	+
	241	+	func TestInteractiveToolHeadless(t *testing.T) {
	242	+	m := New(t.TempDir(), "", "test", "", "", &stubHooks{})
	243	+	tool := NewTool(m, ToolInfo{Name: "ask", Interactive: true})
	244	+	result, err := tool.Execute(context.Background(), nil, nil)
	245	+	if err != nil \|\| !result.IsError \|\| result.Status != "failed" {
	246	+	t.Fatalf("headless result: %+v, %v", result, err)
	247	+	}
	248	+	}
	249	+
	250	+	type blockedToolWriter struct {
	251	+	started chan struct{}
	252	+	closed chan struct{}
	253	+	once sync.Once
	254	+	}
	255	+
	256	+	func (w *blockedToolWriter) Write([]byte) (int, error) {
	257	+	close(w.started)
	258	+	<-w.closed
	259	+	return 0, io.ErrClosedPipe
	260	+	}
	261	+
	262	+	func (w *blockedToolWriter) Close() error {
	263	+	w.once.Do(func() { close(w.closed) })
	264	+	return nil
	265	+	}
	266	+
	267	+	func TestToolCancellationDuringWrite(t *testing.T) {
	268	+	w := &blockedToolWriter{started: make(chan struct{}), closed: make(chan struct{})}
	269	+	pipe := newOrderedPipe(w)
	270	+	defer pipe.Close()
	271	+	ctx, cancel := context.WithCancel(context.Background())
	272	+	defer cancel()
	273	+	done := make(chan error, 1)
	274	+	go func() {
	275	+	_, err := pipe.writeContext(ctx, []byte("{}\n"), time.Hour)
	276	+	done <- err
	277	+	}()
	278	+	select {
	279	+	case <-w.started:
	280	+	case <-time.After(time.Second):
	281	+	t.Fatal("write did not start")
	282	+	}
	283	+	cancel()
	284	+	if err := awaitToolError(t, done); !errors.Is(err, context.Canceled) {
	285	+	t.Fatalf("write cancellation: %v", err)
	286	+	}
	287	+	}

M packages/agent/extensions/manager.go +57 -18


@@ -22,9 +22,11 @@ import (
22	22		"os"
23	23		"os/exec"
24	24		"path/filepath"
	25	+	"slices"
25	26		"sort"
26	27		"strings"
27	28		"sync"
	29	+	"sync/atomic"
28	30		"syscall"
29	31		"time"
30	32

@@ -66,8 +68,10 @@ type Extension struct {
66	68		stdout io.ReadCloser
67	69		logFile *os.File
68	70		helloAck bool
69		-	commands []extproto.RegisterCommandFromExt
70		-	tools []extproto.RegisterToolFromExt
	71	+	// Set during the handshake, before the extension is published.
	72	+	toolCancel bool
	73	+	commands []extproto.RegisterCommandFromExt
	74	+	tools []extproto.RegisterToolFromExt
71	75
72	76		// readyCh is closed when the extension sends a ReadyFromExt
73	77		// frame, or when the host gives up waiting (registrationGrace).

@@ -134,6 +138,12 @@ type HostHooks interface {
134	138		ClosePanel(extName, panelID string)
135	139		}
136	140
	141	+	// Interactive tool support is opt-in so headless hosts fail closed.
	142	+	func (m *Manager) supportsInteractiveTools() bool {
	143	+	h, ok := m.hooks.(interface{ SupportsInteractiveTools() bool })
	144	+	return ok && h.SupportsInteractiveTools()
	145	+	}
	146	+
137	147		type commandRegistration struct {
138	148		ext *Extension
139	149		name string

@@ -768,6 +778,7 @@ func (m Manager) spawn(ctx context.Context, ext Extension) error {
768	778		}
769	779		// Trust the manifest's name; ignore mismatch from the hello.
770	780		ext.helloAck = true
	781	+	ext.toolCancel = slices.Contains(hello.Capabilities, "tool_cancel")
771	782
772	783		ack, _ := extproto.Encode(extproto.HelloAckFromHost{
773	784		Type: "hello_ack",

@@ -864,6 +875,9 @@ func (m Manager) assumeReadyAfterIdle(ext Extension) {
864	875		// Returns when stdout closes.
865	876		func (m Manager) readLoop(ext Extension, scanner *bufio.Scanner) {
866	877		defer func() {
	878	+	if ext.stdin != nil {
	879	+	_ = ext.stdin.Close()
	880	+	}
867	881		// On close, drop every command + tool this extension owned so
868	882		// future invocations don't dangle. The subprocess is gone; we
869	883		// won't hear back about its commands or tool calls anymore.

@@ -1187,6 +1201,7 @@ type ToolInfo struct {
1187	1201		Description string
1188	1202		Schema json.RawMessage
1189	1203		Deferred bool
	1204	+	Interactive bool
1190	1205		}
1191	1206
1192	1207		// Tools returns a snapshot of every (extension, tool) pair currently

@@ -1204,6 +1219,7 @@ func (m *Manager) Tools() []ToolInfo {
1204	1219		Description: t.Description,
1205	1220		Schema: t.Schema,
1206	1221		Deferred: t.Deferred,
	1222	+	Interactive: t.Interactive,
1207	1223		})
1208	1224		}
1209	1225		}

@@ -1222,6 +1238,9 @@ func (m *Manager) HasTool(name string) bool {
1222	1238		// the matching tool_result. Used by the core.Tool wrapper that the
1223	1239		// agent registers per extension-defined tool.
1224	1240		func (m *Manager) InvokeTool(ctx context.Context, name string, args json.RawMessage, timeout time.Duration) (extproto.ToolResultFromExt, error) {
	1241	+	if err := ctx.Err(); err != nil {
	1242	+	return extproto.ToolResultFromExt{}, err
	1243	+	}
1225	1244		m.mu.RLock()
1226	1245		ext, ok := m.toolIndex[name]
1227	1246		m.mu.RUnlock()

@@ -1234,6 +1253,11 @@ func (m *Manager) InvokeTool(ctx context.Context, name string, args json.RawMess
1234	1253		ext.mu.Lock()
1235	1254		ext.pendingTool[id] = ch
1236	1255		ext.mu.Unlock()
	1256	+	defer func() {
	1257	+	ext.mu.Lock()
	1258	+	delete(ext.pendingTool, id)
	1259	+	ext.mu.Unlock()
	1260	+	}()
1237	1261
1238	1262		frame, _ := extproto.Encode(extproto.ToolCallFromHost{
1239	1263		Type: "tool_call",

@@ -1241,26 +1265,41 @@ func (m *Manager) InvokeTool(ctx context.Context, name string, args json.RawMess
1241	1265		Name: name,
1242	1266		Args: args,
1243	1267		})
1244		-	if _, err := ext.stdin.Write(frame); err != nil {
1245		-	ext.mu.Lock()
1246		-	delete(ext.pendingTool, id)
1247		-	ext.mu.Unlock()
	1268	+	var disconnected <-chan struct{}
	1269	+	var err error
	1270	+	pipe, ordered := ext.stdin.(*orderedPipe)
	1271	+	if ordered {
	1272	+	disconnected = pipe.done
	1273	+	_, err = pipe.writeContext(ctx, frame, interceptTimeout)
	1274	+	} else {
	1275	+	_, err = ext.stdin.Write(frame)
	1276	+	}
	1277	+	if err != nil {
1248	1278		return extproto.ToolResultFromExt{}, fmt.Errorf("write: %w", err)
1249	1279		}
1250		-
	1280	+	cancel := func() {
	1281	+	if ordered && ext.toolCancel {
	1282	+	frame, _ := extproto.Encode(extproto.ToolCancelFromHost{Type: "tool_cancel", ID: id})
	1283	+	_ = pipe.enqueue(frame, nil)
	1284	+	}
	1285	+	}
	1286	+	var deadline <-chan time.Time
	1287	+	if timeout > 0 {
	1288	+	timer := time.NewTimer(timeout)
	1289	+	defer timer.Stop()
	1290	+	deadline = timer.C
	1291	+	}
1251	1292		select {
1252	1293		case resp := <-ch:
1253	1294		return resp, nil
1254		-	case <-time.After(timeout):
1255		-	ext.mu.Lock()
1256		-	delete(ext.pendingTool, id)
1257		-	ext.mu.Unlock()
	1295	+	case <-deadline:
	1296	+	cancel()
1258	1297		return extproto.ToolResultFromExt{}, fmt.Errorf("timeout waiting for %s/%s: %w", ext.Manifest.Name, name, context.DeadlineExceeded)
1259	1298		case <-ctx.Done():
1260		-	ext.mu.Lock()
1261		-	delete(ext.pendingTool, id)
1262		-	ext.mu.Unlock()
	1299	+	cancel()
1263	1300		return extproto.ToolResultFromExt{}, ctx.Err()
	1301	+	case <-disconnected:
	1302	+	return extproto.ToolResultFromExt{}, fmt.Errorf("extension %s disconnected", ext.Manifest.Name)
1264	1303		}
1265	1304		}
1266	1305

@@ -1426,9 +1465,9 @@ func (m Manager) All() []Extension {
1426	1465		return out
1427	1466		}
1428	1467
1429		-	// newCorrelationID returns a short non-cryptographic id. We don't
1430		-	// need uniqueness across processes, just within the lifetime of one
1431		-	// extension's pending map.
	1468	+	var correlationSequence atomic.Uint64
	1469	+
	1470	+	// newCorrelationID is unique across invocations in this host process.
1432	1471		func newCorrelationID() string {
1433		-	return strings.ReplaceAll(time.Now().Format("150405.000000"), ".", "")
	1472	+	return fmt.Sprint(correlationSequence.Add(1))
1434	1473		}


@@ -22,9 +22,11 @@ import (
22		"os"	22		"os"
23		"os/exec"	23		"os/exec"
24		"path/filepath"	24		"path/filepath"
			25	+	"slices"
25		"sort"	26		"sort"
26		"strings"	27		"strings"
27		"sync"	28		"sync"
			29	+	"sync/atomic"
28		"syscall"	30		"syscall"
29		"time"	31		"time"
30			32

@@ -66,8 +68,10 @@ type Extension struct {
66		stdout io.ReadCloser	68		stdout io.ReadCloser
67		logFile *os.File	69		logFile *os.File
68		helloAck bool	70		helloAck bool
69	-	commands []extproto.RegisterCommandFromExt	71	+	// Set during the handshake, before the extension is published.
70	-	tools []extproto.RegisterToolFromExt	72	+	toolCancel bool
			73	+	commands []extproto.RegisterCommandFromExt
			74	+	tools []extproto.RegisterToolFromExt
71			75
72		// readyCh is closed when the extension sends a ReadyFromExt	76		// readyCh is closed when the extension sends a ReadyFromExt
73		// frame, or when the host gives up waiting (registrationGrace).	77		// frame, or when the host gives up waiting (registrationGrace).

@@ -134,6 +138,12 @@ type HostHooks interface {
134		ClosePanel(extName, panelID string)	138		ClosePanel(extName, panelID string)
135		}	139		}
136			140
			141	+	// Interactive tool support is opt-in so headless hosts fail closed.
			142	+	func (m *Manager) supportsInteractiveTools() bool {
			143	+	h, ok := m.hooks.(interface{ SupportsInteractiveTools() bool })
			144	+	return ok && h.SupportsInteractiveTools()
			145	+	}
			146	+
137		type commandRegistration struct {	147		type commandRegistration struct {
138		ext *Extension	148		ext *Extension
139		name string	149		name string

@@ -768,6 +778,7 @@ func (m Manager) spawn(ctx context.Context, ext Extension) error {
768		}	778		}
769		// Trust the manifest's name; ignore mismatch from the hello.	779		// Trust the manifest's name; ignore mismatch from the hello.
770		ext.helloAck = true	780		ext.helloAck = true
			781	+	ext.toolCancel = slices.Contains(hello.Capabilities, "tool_cancel")
771			782
772		ack, _ := extproto.Encode(extproto.HelloAckFromHost{	783		ack, _ := extproto.Encode(extproto.HelloAckFromHost{
773		Type: "hello_ack",	784		Type: "hello_ack",

@@ -864,6 +875,9 @@ func (m Manager) assumeReadyAfterIdle(ext Extension) {
864		// Returns when stdout closes.	875		// Returns when stdout closes.
865		func (m Manager) readLoop(ext Extension, scanner *bufio.Scanner) {	876		func (m Manager) readLoop(ext Extension, scanner *bufio.Scanner) {
866		defer func() {	877		defer func() {
			878	+	if ext.stdin != nil {
			879	+	_ = ext.stdin.Close()
			880	+	}
867		// On close, drop every command + tool this extension owned so	881		// On close, drop every command + tool this extension owned so
868		// future invocations don't dangle. The subprocess is gone; we	882		// future invocations don't dangle. The subprocess is gone; we
869		// won't hear back about its commands or tool calls anymore.	883		// won't hear back about its commands or tool calls anymore.

@@ -1187,6 +1201,7 @@ type ToolInfo struct {
1187		Description string	1201		Description string
1188		Schema json.RawMessage	1202		Schema json.RawMessage
1189		Deferred bool	1203		Deferred bool
			1204	+	Interactive bool
1190		}	1205		}
1191			1206
1192		// Tools returns a snapshot of every (extension, tool) pair currently	1207		// Tools returns a snapshot of every (extension, tool) pair currently

@@ -1204,6 +1219,7 @@ func (m *Manager) Tools() []ToolInfo {
1204		Description: t.Description,	1219		Description: t.Description,
1205		Schema: t.Schema,	1220		Schema: t.Schema,
1206		Deferred: t.Deferred,	1221		Deferred: t.Deferred,
			1222	+	Interactive: t.Interactive,
1207		})	1223		})
1208		}	1224		}
1209		}	1225		}

@@ -1222,6 +1238,9 @@ func (m *Manager) HasTool(name string) bool {
1222		// the matching tool_result. Used by the core.Tool wrapper that the	1238		// the matching tool_result. Used by the core.Tool wrapper that the
1223		// agent registers per extension-defined tool.	1239		// agent registers per extension-defined tool.
1224		func (m *Manager) InvokeTool(ctx context.Context, name string, args json.RawMessage, timeout time.Duration) (extproto.ToolResultFromExt, error) {	1240		func (m *Manager) InvokeTool(ctx context.Context, name string, args json.RawMessage, timeout time.Duration) (extproto.ToolResultFromExt, error) {
			1241	+	if err := ctx.Err(); err != nil {
			1242	+	return extproto.ToolResultFromExt{}, err
			1243	+	}
1225		m.mu.RLock()	1244		m.mu.RLock()
1226		ext, ok := m.toolIndex[name]	1245		ext, ok := m.toolIndex[name]
1227		m.mu.RUnlock()	1246		m.mu.RUnlock()

@@ -1234,6 +1253,11 @@ func (m *Manager) InvokeTool(ctx context.Context, name string, args json.RawMess
1234		ext.mu.Lock()	1253		ext.mu.Lock()
1235		ext.pendingTool[id] = ch	1254		ext.pendingTool[id] = ch
1236		ext.mu.Unlock()	1255		ext.mu.Unlock()
			1256	+	defer func() {
			1257	+	ext.mu.Lock()
			1258	+	delete(ext.pendingTool, id)
			1259	+	ext.mu.Unlock()
			1260	+	}()
1237			1261
1238		frame, _ := extproto.Encode(extproto.ToolCallFromHost{	1262		frame, _ := extproto.Encode(extproto.ToolCallFromHost{
1239		Type: "tool_call",	1263		Type: "tool_call",

@@ -1241,26 +1265,41 @@ func (m *Manager) InvokeTool(ctx context.Context, name string, args json.RawMess
1241		Name: name,	1265		Name: name,
1242		Args: args,	1266		Args: args,
1243		})	1267		})
1244	-	if _, err := ext.stdin.Write(frame); err != nil {	1268	+	var disconnected <-chan struct{}
1245	-	ext.mu.Lock()	1269	+	var err error
1246	-	delete(ext.pendingTool, id)	1270	+	pipe, ordered := ext.stdin.(*orderedPipe)
1247	-	ext.mu.Unlock()	1271	+	if ordered {
			1272	+	disconnected = pipe.done
			1273	+	_, err = pipe.writeContext(ctx, frame, interceptTimeout)
			1274	+	} else {
			1275	+	_, err = ext.stdin.Write(frame)
			1276	+	}
			1277	+	if err != nil {
1248		return extproto.ToolResultFromExt{}, fmt.Errorf("write: %w", err)	1278		return extproto.ToolResultFromExt{}, fmt.Errorf("write: %w", err)
1249		}	1279		}
1250	-		1280	+	cancel := func() {
			1281	+	if ordered && ext.toolCancel {
			1282	+	frame, _ := extproto.Encode(extproto.ToolCancelFromHost{Type: "tool_cancel", ID: id})
			1283	+	_ = pipe.enqueue(frame, nil)
			1284	+	}
			1285	+	}
			1286	+	var deadline <-chan time.Time
			1287	+	if timeout > 0 {
			1288	+	timer := time.NewTimer(timeout)
			1289	+	defer timer.Stop()
			1290	+	deadline = timer.C
			1291	+	}
1251		select {	1292		select {
1252		case resp := <-ch:	1293		case resp := <-ch:
1253		return resp, nil	1294		return resp, nil
1254	-	case <-time.After(timeout):	1295	+	case <-deadline:
1255	-	ext.mu.Lock()	1296	+	cancel()
1256	-	delete(ext.pendingTool, id)
1257	-	ext.mu.Unlock()
1258		return extproto.ToolResultFromExt{}, fmt.Errorf("timeout waiting for %s/%s: %w", ext.Manifest.Name, name, context.DeadlineExceeded)	1297		return extproto.ToolResultFromExt{}, fmt.Errorf("timeout waiting for %s/%s: %w", ext.Manifest.Name, name, context.DeadlineExceeded)
1259		case <-ctx.Done():	1298		case <-ctx.Done():
1260	-	ext.mu.Lock()	1299	+	cancel()
1261	-	delete(ext.pendingTool, id)
1262	-	ext.mu.Unlock()
1263		return extproto.ToolResultFromExt{}, ctx.Err()	1300		return extproto.ToolResultFromExt{}, ctx.Err()
			1301	+	case <-disconnected:
			1302	+	return extproto.ToolResultFromExt{}, fmt.Errorf("extension %s disconnected", ext.Manifest.Name)
1264		}	1303		}
1265		}	1304		}
1266			1305

@@ -1426,9 +1465,9 @@ func (m Manager) All() []Extension {
1426		return out	1465		return out
1427		}	1466		}
1428			1467
1429	-	// newCorrelationID returns a short non-cryptographic id. We don't	1468	+	var correlationSequence atomic.Uint64
1430	-	// need uniqueness across processes, just within the lifetime of one	1469	+
1431	-	// extension's pending map.	1470	+	// newCorrelationID is unique across invocations in this host process.
1432		func newCorrelationID() string {	1471		func newCorrelationID() string {
1433	-	return strings.ReplaceAll(time.Now().Format("150405.000000"), ".", "")	1472	+	return fmt.Sprint(correlationSequence.Add(1))
1434		}	1473		}

M packages/agent/extensions/outbound.go +11 -0


@@ -65,6 +65,13 @@ func (p *orderedPipe) Write(data []byte) (int, error) {
65	65		}
66	66
67	67		func (p *orderedPipe) writeTimeout(data []byte, timeout time.Duration) (int, error) {
	68	+	return p.writeContext(context.Background(), data, timeout)
	69	+	}
	70	+
	71	+	func (p *orderedPipe) writeContext(ctx context.Context, data []byte, timeout time.Duration) (int, error) {
	72	+	if err := ctx.Err(); err != nil {
	73	+	return 0, err
	74	+	}
68	75		if timeout <= 0 {
69	76		p.Close()
70	77		return 0, io.ErrClosedPipe

@@ -83,6 +90,10 @@ func (p *orderedPipe) writeTimeout(data []byte, timeout time.Duration) (int, err
83	90		return len(data), nil
84	91		case <-p.done:
85	92		return 0, io.ErrClosedPipe
	93	+	case <-ctx.Done():
	94	+	// A partially written frame cannot be safely withdrawn.
	95	+	p.Close()
	96	+	return 0, ctx.Err()
86	97		case <-timer.C:
87	98		p.Close()
88	99		return 0, fmt.Errorf("extension outbound write timed out: %w", context.DeadlineExceeded)


@@ -65,6 +65,13 @@ func (p *orderedPipe) Write(data []byte) (int, error) {
65	}	65		}
66		66
67	func (p *orderedPipe) writeTimeout(data []byte, timeout time.Duration) (int, error) {	67		func (p *orderedPipe) writeTimeout(data []byte, timeout time.Duration) (int, error) {
		68	+	return p.writeContext(context.Background(), data, timeout)
		69	+	}
		70	+
		71	+	func (p *orderedPipe) writeContext(ctx context.Context, data []byte, timeout time.Duration) (int, error) {
		72	+	if err := ctx.Err(); err != nil {
		73	+	return 0, err
		74	+	}
68	if timeout <= 0 {	75		if timeout <= 0 {
69	p.Close()	76		p.Close()
70	return 0, io.ErrClosedPipe	77		return 0, io.ErrClosedPipe

@@ -83,6 +90,10 @@ func (p *orderedPipe) writeTimeout(data []byte, timeout time.Duration) (int, err
83	return len(data), nil	90		return len(data), nil
84	case <-p.done:	91		case <-p.done:
85	return 0, io.ErrClosedPipe	92		return 0, io.ErrClosedPipe
		93	+	case <-ctx.Done():
		94	+	// A partially written frame cannot be safely withdrawn.
		95	+	p.Close()
		96	+	return 0, ctx.Err()
86	case <-timer.C:	97		case <-timer.C:
87	p.Close()	98		p.Close()
88	return 0, fmt.Errorf("extension outbound write timed out: %w", context.DeadlineExceeded)	99		return 0, fmt.Errorf("extension outbound write timed out: %w", context.DeadlineExceeded)

M packages/agent/extensions/tool.go +13 -2


@@ -27,20 +27,26 @@ type extensionTool struct {
27	27		manager *Manager
28	28		timeout time.Duration
29	29		deferred bool
	30	+	interactive bool
30	31		}
31	32
32	33		// NewTool returns a core.Tool that round-trips invocations through
33	34		// mgr to the extension that registered (name, schema). The default
34		-	// per-call timeout is 60 seconds; callers can override.
	35	+	// per-call timeout is 60 seconds. Interactive tools have no reply deadline.
35	36		func NewTool(mgr *Manager, info ToolInfo) core.Tool {
	37	+	timeout := 60 * time.Second
	38	+	if info.Interactive {
	39	+	timeout = 0
	40	+	}
36	41		return &extensionTool{
37	42		name: info.Name,
38	43		description: info.Description,
39	44		schema: info.Schema,
40	45		extension: info.Extension,
41	46		manager: mgr,
42		-	timeout: 60 * time.Second,
	47	+	timeout: timeout,
43	48		deferred: info.Deferred,
	49	+	interactive: info.Interactive,
44	50		}
45	51		}
46	52

@@ -57,6 +63,11 @@ func (t *extensionTool) Execute(ctx context.Context, args json.RawMessage, _ fun
57	63		if len(args) == 0 {
58	64		args = json.RawMessage(`{}`)
59	65		}
	66	+	if t.interactive && !t.manager.supportsInteractiveTools() {
	67	+	return core.ToolResult{Status: "failed", IsError: true, Content: []provider.Content{
	68	+	provider.TextBlock{Text: "interactive extension tools require an interactive host"},
	69	+	}}, nil
	70	+	}
60	71		resp, err := t.manager.InvokeTool(ctx, t.name, args, t.timeout)
61	72		if err != nil {
62	73		status := "failed"

 	manager     *Manager
 	timeout     time.Duration
 	deferred    bool
+	interactive bool
 }
 // NewTool returns a core.Tool that round-trips invocations through
 // mgr to the extension that registered (name, schema). The default
-// per-call timeout is 60 seconds; callers can override.
+// per-call timeout is 60 seconds. Interactive tools have no reply deadline.
 func NewTool(mgr *Manager, info ToolInfo) core.Tool {
+	timeout := 60 * time.Second
+	if info.Interactive {
+		timeout = 0
+	}
 	return &extensionTool{
 		name:        info.Name,
 		description: info.Description,
 		schema:      info.Schema,
 		extension:   info.Extension,
 		manager:     mgr,
-		timeout:     60 * time.Second,
+		timeout:     timeout,
 		deferred:    info.Deferred,
+		interactive: info.Interactive,
 	}
 }
 	if len(args) == 0 {
 		args = json.RawMessage(`{}`)
 	}
+	if t.interactive && !t.manager.supportsInteractiveTools() {
+		return core.ToolResult{Status: "failed", IsError: true, Content: []provider.Content{
+			provider.TextBlock{Text: "interactive extension tools require an interactive host"},
+		}}, nil
+	}
 	resp, err := t.manager.InvokeTool(ctx, t.name, args, t.timeout)
 	if err != nil {
 		status := "failed"

A packages/agent/extensions/tool_cancel_test.go +137 -0

@@ -0,0 +1,137 @@
	1	+	package extensions
	2	+
	3	+	import (
	4	+	"context"
	5	+	"encoding/json"
	6	+	"errors"
	7	+	"os"
	8	+	"path/filepath"
	9	+	"strings"
	10	+	"testing"
	11	+	"testing/synctest"
	12	+	"time"
	13	+
	14	+	"github.com/patriceckhart/zot/packages/agent/extproto"
	15	+	)
	16	+
	17	+	func TestToolCancelCapabilityHandshake(t *testing.T) {
	18	+	for _, tc := range []struct {
	19	+	name string
	20	+	caps []string
	21	+	want bool
	22	+	}{
	23	+	{name: "omitted"},
	24	+	{name: "legacy", caps: []string{"tools", "panels"}},
	25	+	{name: "unknown", caps: []string{"tool_cancel_future", "TOOL_CANCEL"}},
	26	+	{name: "supported", caps: []string{"tools", "tool_cancel"}, want: true},
	27	+	} {
	28	+	t.Run(tc.name, func(t *testing.T) {
	29	+	root := t.TempDir()
	30	+	writeMockToolExtension(t, filepath.Join(root, "extensions"))
	31	+	path := filepath.Join(root, "extensions", "tool-mock", "run.sh")
	32	+	script, err := os.ReadFile(path)
	33	+	if err != nil {
	34	+	t.Fatal(err)
	35	+	}
	36	+	hello, err := json.Marshal(extproto.HelloFromExt{Type: "hello", Name: "tool-mock", Version: "0.1", Capabilities: tc.caps})
	37	+	if err != nil {
	38	+	t.Fatal(err)
	39	+	}
	40	+	lines := strings.Split(string(script), "\n")
	41	+	lines[1] = "printf '%s\\n' '" + string(hello) + "'"
	42	+	if err := os.WriteFile(path, []byte(strings.Join(lines, "\n")), 0o755); err != nil {
	43	+	t.Fatal(err)
	44	+	}
	45	+	m := New(root, "", "test", "", "", &stubHooks{})
	46	+	t.Cleanup(func() { m.Stop(time.Second) })
	47	+	if errs := m.Discover(context.Background()); len(errs) != 0 {
	48	+	t.Fatal(errs)
	49	+	}
	50	+	m.WaitForReady(time.Second)
	51	+	m.mu.RLock()
	52	+	ext := m.toolIndex["echo"]
	53	+	m.mu.RUnlock()
	54	+	if ext == nil {
	55	+	t.Fatal("tool was not registered")
	56	+	}
	57	+	if ext.toolCancel != tc.want {
	58	+	t.Fatalf("cancellation support = %v, want %v", ext.toolCancel, tc.want)
	59	+	}
	60	+	})
	61	+	}
	62	+	}
	63	+
	64	+	func TestToolCancelCapabilityDelivery(t *testing.T) {
	65	+	for _, supported := range []bool{false, true} {
	66	+	name := "legacy"
	67	+	if supported {
	68	+	name = "supported"
	69	+	}
	70	+	for _, outcome := range []string{"cancel", "timeout", "complete"} {
	71	+	t.Run(name+"/"+outcome, func(t *testing.T) {
	72	+	synctest.Test(t, func(t *testing.T) {
	73	+	m, ext, frames, replies := toolPeer(t)
	74	+	ext.toolCancel = supported
	75	+	ctx, cancel := context.WithCancel(context.Background())
	76	+	defer cancel()
	77	+	timeout := time.Duration(0)
	78	+	if outcome == "timeout" {
	79	+	timeout = time.Nanosecond
	80	+	}
	81	+	done := make(chan error, 1)
	82	+	go func() {
	83	+	_, err := m.InvokeTool(ctx, "ask", nil, timeout)
	84	+	done <- err
	85	+	}()
	86	+	call := toolCallFrame(t, frames)
	87	+	synctest.Wait()
	88	+	var want error
	89	+	switch outcome {
	90	+	case "cancel":
	91	+	want = context.Canceled
	92	+	cancel()
	93	+	case "timeout":
	94	+	want = context.DeadlineExceeded
	95	+	case "complete":
	96	+	sendToolFrame(t, replies, extproto.ToolResultFromExt{Type: "tool_result", ID: call.ID})
	97	+	}
	98	+	if err := awaitToolError(t, done); !errors.Is(err, want) {
	99	+	t.Fatalf("invocation error = %v, want %v", err, want)
	100	+	}
	101	+	if supported && outcome != "complete" {
	102	+	var frame extproto.ToolCancelFromHost
	103	+	if err := json.Unmarshal(nextLifecycleFrame(t, frames), &frame); err != nil {
	104	+	t.Fatal(err)
	105	+	}
	106	+	if frame.Type != "tool_cancel" \|\| frame.ID != call.ID {
	107	+	t.Fatalf("unexpected cancellation frame: %+v", frame)
	108	+	}
	109	+	}
	110	+	synctest.Wait()
	111	+	select {
	112	+	case frame := <-frames:
	113	+	t.Fatalf("unexpected host frame: %s", frame)
	114	+	default:
	115	+	}
	116	+	ext.mu.Lock()
	117	+	pending := len(ext.pendingTool)
	118	+	ext.mu.Unlock()
	119	+	if pending != 0 {
	120	+	t.Fatal("pending invocation leaked")
	121	+	}
	122	+	// Abandoned results remain harmless without cancellation support.
	123	+	sendToolFrame(t, replies, extproto.ToolResultFromExt{Type: "tool_result", ID: call.ID})
	124	+	go func() {
	125	+	_, err := m.InvokeTool(context.Background(), "ask", nil, 0)
	126	+	done <- err
	127	+	}()
	128	+	next := toolCallFrame(t, frames)
	129	+	sendToolFrame(t, replies, extproto.ToolResultFromExt{Type: "tool_result", ID: next.ID})
	130	+	if err := awaitToolError(t, done); err != nil {
	131	+	t.Fatal(err)
	132	+	}
	133	+	})
	134	+	})
	135	+	}
	136	+	}
	137	+	}

@@ -0,0 +1,137 @@
	1	+	package extensions
	2	+
	3	+	import (
	4	+	"context"
	5	+	"encoding/json"
	6	+	"errors"
	7	+	"os"
	8	+	"path/filepath"
	9	+	"strings"
	10	+	"testing"
	11	+	"testing/synctest"
	12	+	"time"
	13	+
	14	+	"github.com/patriceckhart/zot/packages/agent/extproto"
	15	+	)
	16	+
	17	+	func TestToolCancelCapabilityHandshake(t *testing.T) {
	18	+	for _, tc := range []struct {
	19	+	name string
	20	+	caps []string
	21	+	want bool
	22	+	}{
	23	+	{name: "omitted"},
	24	+	{name: "legacy", caps: []string{"tools", "panels"}},
	25	+	{name: "unknown", caps: []string{"tool_cancel_future", "TOOL_CANCEL"}},
	26	+	{name: "supported", caps: []string{"tools", "tool_cancel"}, want: true},
	27	+	} {
	28	+	t.Run(tc.name, func(t *testing.T) {
	29	+	root := t.TempDir()
	30	+	writeMockToolExtension(t, filepath.Join(root, "extensions"))
	31	+	path := filepath.Join(root, "extensions", "tool-mock", "run.sh")
	32	+	script, err := os.ReadFile(path)
	33	+	if err != nil {
	34	+	t.Fatal(err)
	35	+	}
	36	+	hello, err := json.Marshal(extproto.HelloFromExt{Type: "hello", Name: "tool-mock", Version: "0.1", Capabilities: tc.caps})
	37	+	if err != nil {
	38	+	t.Fatal(err)
	39	+	}
	40	+	lines := strings.Split(string(script), "\n")
	41	+	lines[1] = "printf '%s\\n' '" + string(hello) + "'"
	42	+	if err := os.WriteFile(path, []byte(strings.Join(lines, "\n")), 0o755); err != nil {
	43	+	t.Fatal(err)
	44	+	}
	45	+	m := New(root, "", "test", "", "", &stubHooks{})
	46	+	t.Cleanup(func() { m.Stop(time.Second) })
	47	+	if errs := m.Discover(context.Background()); len(errs) != 0 {
	48	+	t.Fatal(errs)
	49	+	}
	50	+	m.WaitForReady(time.Second)
	51	+	m.mu.RLock()
	52	+	ext := m.toolIndex["echo"]
	53	+	m.mu.RUnlock()
	54	+	if ext == nil {
	55	+	t.Fatal("tool was not registered")
	56	+	}
	57	+	if ext.toolCancel != tc.want {
	58	+	t.Fatalf("cancellation support = %v, want %v", ext.toolCancel, tc.want)
	59	+	}
	60	+	})
	61	+	}
	62	+	}
	63	+
	64	+	func TestToolCancelCapabilityDelivery(t *testing.T) {
	65	+	for _, supported := range []bool{false, true} {
	66	+	name := "legacy"
	67	+	if supported {
	68	+	name = "supported"
	69	+	}
	70	+	for _, outcome := range []string{"cancel", "timeout", "complete"} {
	71	+	t.Run(name+"/"+outcome, func(t *testing.T) {
	72	+	synctest.Test(t, func(t *testing.T) {
	73	+	m, ext, frames, replies := toolPeer(t)
	74	+	ext.toolCancel = supported
	75	+	ctx, cancel := context.WithCancel(context.Background())
	76	+	defer cancel()
	77	+	timeout := time.Duration(0)
	78	+	if outcome == "timeout" {
	79	+	timeout = time.Nanosecond
	80	+	}
	81	+	done := make(chan error, 1)
	82	+	go func() {
	83	+	_, err := m.InvokeTool(ctx, "ask", nil, timeout)
	84	+	done <- err
	85	+	}()
	86	+	call := toolCallFrame(t, frames)
	87	+	synctest.Wait()
	88	+	var want error
	89	+	switch outcome {
	90	+	case "cancel":
	91	+	want = context.Canceled
	92	+	cancel()
	93	+	case "timeout":
	94	+	want = context.DeadlineExceeded
	95	+	case "complete":
	96	+	sendToolFrame(t, replies, extproto.ToolResultFromExt{Type: "tool_result", ID: call.ID})
	97	+	}
	98	+	if err := awaitToolError(t, done); !errors.Is(err, want) {
	99	+	t.Fatalf("invocation error = %v, want %v", err, want)
	100	+	}
	101	+	if supported && outcome != "complete" {
	102	+	var frame extproto.ToolCancelFromHost
	103	+	if err := json.Unmarshal(nextLifecycleFrame(t, frames), &frame); err != nil {
	104	+	t.Fatal(err)
	105	+	}
	106	+	if frame.Type != "tool_cancel" \|\| frame.ID != call.ID {
	107	+	t.Fatalf("unexpected cancellation frame: %+v", frame)
	108	+	}
	109	+	}
	110	+	synctest.Wait()
	111	+	select {
	112	+	case frame := <-frames:
	113	+	t.Fatalf("unexpected host frame: %s", frame)
	114	+	default:
	115	+	}
	116	+	ext.mu.Lock()
	117	+	pending := len(ext.pendingTool)
	118	+	ext.mu.Unlock()
	119	+	if pending != 0 {
	120	+	t.Fatal("pending invocation leaked")
	121	+	}
	122	+	// Abandoned results remain harmless without cancellation support.
	123	+	sendToolFrame(t, replies, extproto.ToolResultFromExt{Type: "tool_result", ID: call.ID})
	124	+	go func() {
	125	+	_, err := m.InvokeTool(context.Background(), "ask", nil, 0)
	126	+	done <- err
	127	+	}()
	128	+	next := toolCallFrame(t, frames)
	129	+	sendToolFrame(t, replies, extproto.ToolResultFromExt{Type: "tool_result", ID: next.ID})
	130	+	if err := awaitToolError(t, done); err != nil {
	131	+	t.Fatal(err)
	132	+	}
	133	+	})
	134	+	})
	135	+	}
	136	+	}
	137	+	}

M packages/agent/extproto/extproto.go +7 -0


@@ -45,6 +45,7 @@ type RegisterToolFromExt struct {
45	45		Description string `json:"description,omitempty"`
46	46		Schema json.RawMessage `json:"schema"`
47	47		Deferred bool `json:"deferred,omitempty"`
	48	+	Interactive bool `json:"interactive,omitempty"`
48	49		}
49	50
50	51		type ReadyFromExt struct {

@@ -176,6 +177,12 @@ type CommandInvokedFromHost struct {
176	177		Args string `json:"args,omitempty"`
177	178		}
178	179
	180	+	// ToolCancelFromHost asks the extension to stop an abandoned invocation.
	181	+	type ToolCancelFromHost struct {
	182	+	Type string `json:"type"`
	183	+	ID string `json:"id"`
	184	+	}
	185	+
179	186		type ToolCallFromHost struct {
180	187		Type string `json:"type"`
181	188		ID string `json:"id"`


@@ -45,6 +45,7 @@ type RegisterToolFromExt struct {
45	Description string `json:"description,omitempty"`	45		Description string `json:"description,omitempty"`
46	Schema json.RawMessage `json:"schema"`	46		Schema json.RawMessage `json:"schema"`
47	Deferred bool `json:"deferred,omitempty"`	47		Deferred bool `json:"deferred,omitempty"`
		48	+	Interactive bool `json:"interactive,omitempty"`
48	}	49		}
49		50
50	type ReadyFromExt struct {	51		type ReadyFromExt struct {

@@ -176,6 +177,12 @@ type CommandInvokedFromHost struct {
176	Args string `json:"args,omitempty"`	177		Args string `json:"args,omitempty"`
177	}	178		}
178		179
		180	+	// ToolCancelFromHost asks the extension to stop an abandoned invocation.
		181	+	type ToolCancelFromHost struct {
		182	+	Type string `json:"type"`
		183	+	ID string `json:"id"`
		184	+	}
		185	+
179	type ToolCallFromHost struct {	186		type ToolCallFromHost struct {
180	Type string `json:"type"`	187		Type string `json:"type"`
181	ID string `json:"id"`	188		ID string `json:"id"`