1package extensions
2
3import (
4 "context"
5 "encoding/json"
6 "errors"
7 "fmt"
8 "time"
9
10 "github.com/bnema/zut/packages/core"
11 "github.com/bnema/zut/packages/provider"
12)
13
14// extensionTool wraps a single extension-registered tool as a
15// core.Tool. The agent's tool registry contains one of these per
16// extension tool; Execute round-trips through the manager to the
17// owning subprocess.
18//
19// One concrete type instead of a closure-driven anonymous tool
20// keeps the schema, name, and ownership inspectable for logs and
21// dialogs.
22// ToolResultDetails identifies extension metadata carried out of the
23// subprocess. State is persisted with the active session when present and
24// is never serialized into a provider request.
25type ToolResultDetails struct {
26 Extension string
27 Tool string
28 State json.RawMessage
29}
30
31type extensionTool struct {
32 name string
33 description string
34 schema json.RawMessage
35 extension string
36 manager *Manager
37 timeout time.Duration
38 deferred bool
39 interactive bool
40}
41
42// defaultToolTimeout is the reply deadline for ordinary extension tools.
43const defaultToolTimeout = 60 * time.Second
44
45// InteractiveTool is implemented by extension tools that wait for user
46// input. Hosts that build child or headless registries from a parent's
47// catalogue can use IsInteractiveTool to exclude them explicitly instead of
48// relying on the runtime rejection in Execute.
49type InteractiveTool interface {
50 Interactive() bool
51}
52
53// IsInteractiveTool reports whether t is an interactive extension tool.
54func IsInteractiveTool(t core.Tool) bool {
55 it, ok := t.(InteractiveTool)
56 return ok && it.Interactive()
57}
58
59// NewTool returns a core.Tool that round-trips invocations through
60// mgr to the extension that registered (name, schema). Ordinary tools
61// must reply within 60 seconds. Interactive tools have no host reply
62// deadline; the agent context, its deadline, and extension disconnects
63// still end the call.
64func NewTool(mgr *Manager, info ToolInfo) core.Tool {
65 timeout := defaultToolTimeout
66 if info.Interactive {
67 timeout = 0
68 }
69 return &extensionTool{
70 name: info.Name,
71 description: info.Description,
72 schema: info.Schema,
73 extension: info.Extension,
74 manager: mgr,
75 timeout: timeout,
76 deferred: info.Deferred,
77 interactive: info.Interactive,
78 }
79}
80
81func (t *extensionTool) Name() string { return t.name }
82func (t *extensionTool) Description() string { return t.description }
83func (t *extensionTool) Schema() json.RawMessage { return t.schema }
84func (t *extensionTool) Extension() string { return t.extension }
85func (t *extensionTool) Deferred() bool { return t.deferred }
86func (t *extensionTool) Interactive() bool { return t.interactive }
87
88// Execute is what the agent calls when the LLM invokes the tool. It
89// hands args to the owning extension, waits up to t.timeout for the
90// reply, and converts the response into a core.ToolResult.
91func (t *extensionTool) Execute(ctx context.Context, args json.RawMessage, _ func(string)) (core.ToolResult, error) {
92 if len(args) == 0 {
93 args = json.RawMessage(`{}`)
94 }
95 if t.interactive && !t.manager.supportsInteractiveTools() {
96 // Fail closed: a headless host has nowhere to show the prompt, so
97 // waiting without a deadline would hang the run.
98 return core.ToolResult{
99 IsError: true,
100 Content: []provider.Content{provider.TextBlock{Text: fmt.Sprintf("extension %s/%s is interactive and requires an interactive host", t.extension, t.name)}},
101 }, nil
102 }
103 resp, err := t.manager.InvokeTool(ctx, t.name, args, t.timeout)
104 if err != nil {
105 if errors.Is(err, context.Canceled) {
106 return core.ToolResult{
107 IsError: true,
108 Content: []provider.Content{provider.TextBlock{Text: fmt.Sprintf("extension %s/%s cancelled: %v", t.extension, t.name, err)}},
109 }, nil
110 }
111 return core.ToolResult{
112 IsError: true,
113 Content: []provider.Content{provider.TextBlock{Text: fmt.Sprintf("extension %s/%s failed: %v", t.extension, t.name, err)}},
114 }, nil
115 }
116 out := core.ToolResult{IsError: resp.IsError, ActivateTools: resp.ActivateTools}
117 for _, b := range resp.Content {
118 switch b.Type {
119 case "text":
120 if b.Text != "" {
121 out.Content = append(out.Content, provider.TextBlock{Text: b.Text})
122 }
123 case "image":
124 data, dErr := decodeBase64(b.Data)
125 if dErr != nil {
126 out.IsError = true
127 out.Content = append(out.Content, provider.TextBlock{Text: fmt.Sprintf("extension %s/%s returned undecodable image: %v", t.extension, t.name, dErr)})
128 continue
129 }
130 out.Content = append(out.Content, provider.ImageBlock{
131 MimeType: b.MimeType,
132 Data: data,
133 })
134 default:
135 // Unknown block type: stringify for debug visibility.
136 out.Content = append(out.Content, provider.TextBlock{Text: fmt.Sprintf("[unknown block type %q from extension %s]", b.Type, t.extension)})
137 }
138 }
139 if len(out.Content) == 0 {
140 // Defensive: an empty content slice would confuse the model.
141 out.Content = []provider.Content{provider.TextBlock{Text: "(extension returned no content)"}}
142 }
143 out.Details = ToolResultDetails{
144 Extension: t.extension,
145 Tool: t.name,
146 State: append(json.RawMessage(nil), resp.Details...),
147 }
148 return out, nil
149}
150
151// decodeBase64 is a tiny wrapper around encoding/base64 so we can
152// validate the extension's image data in one place.
153func decodeBase64(s string) ([]byte, error) {
154 if s == "" {
155 return nil, nil
156 }
157 return base64DecodeStd(s)
158}