1package extensions
2
3import (
4 "context"
5 "encoding/json"
6 "errors"
7 "fmt"
8 "time"
9
10 "github.com/patriceckhart/zot/packages/core"
11 "github.com/patriceckhart/zot/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.
22type extensionTool struct {
23 name string
24 description string
25 schema json.RawMessage
26 extension string
27 manager *Manager
28 timeout time.Duration
29 deferred bool
30}
31
32// NewTool returns a core.Tool that round-trips invocations through
33// mgr to the extension that registered (name, schema). The default
34// per-call timeout is 60 seconds; callers can override.
35func NewTool(mgr *Manager, info ToolInfo) core.Tool {
36 return &extensionTool{
37 name: info.Name,
38 description: info.Description,
39 schema: info.Schema,
40 extension: info.Extension,
41 manager: mgr,
42 timeout: 60 * time.Second,
43 deferred: info.Deferred,
44 }
45}
46
47func (t *extensionTool) Name() string { return t.name }
48func (t *extensionTool) Description() string { return t.description }
49func (t *extensionTool) Schema() json.RawMessage { return t.schema }
50func (t *extensionTool) Extension() string { return t.extension }
51func (t *extensionTool) Deferred() bool { return t.deferred }
52
53// Execute is what the agent calls when the LLM invokes the tool. It
54// hands args to the owning extension, waits up to t.timeout for the
55// reply, and converts the response into a core.ToolResult.
56func (t *extensionTool) Execute(ctx context.Context, args json.RawMessage, _ func(string)) (core.ToolResult, error) {
57 if len(args) == 0 {
58 args = json.RawMessage(`{}`)
59 }
60 resp, err := t.manager.InvokeTool(ctx, t.name, args, t.timeout)
61 if err != nil {
62 status := "failed"
63 if errors.Is(err, context.Canceled) {
64 status = "cancelled"
65 }
66 if errors.Is(err, context.DeadlineExceeded) {
67 status = "timed_out"
68 }
69 return core.ToolResult{
70 Status: status,
71 IsError: true,
72 Content: []provider.Content{provider.TextBlock{Text: fmt.Sprintf("extension %s/%s failed: %v", t.extension, t.name, err)}},
73 }, nil
74 }
75 out := core.ToolResult{IsError: resp.IsError, ActivateTools: resp.ActivateTools}
76 for _, b := range resp.Content {
77 switch b.Type {
78 case "text":
79 if b.Text != "" {
80 out.Content = append(out.Content, provider.TextBlock{Text: b.Text})
81 }
82 case "image":
83 data, dErr := decodeBase64(b.Data)
84 if dErr != nil {
85 out.IsError = true
86 out.Content = append(out.Content, provider.TextBlock{Text: fmt.Sprintf("extension %s/%s returned undecodable image: %v", t.extension, t.name, dErr)})
87 continue
88 }
89 out.Content = append(out.Content, provider.ImageBlock{
90 MimeType: b.MimeType,
91 Data: data,
92 })
93 default:
94 // Unknown block type: stringify for debug visibility.
95 out.Content = append(out.Content, provider.TextBlock{Text: fmt.Sprintf("[unknown block type %q from extension %s]", b.Type, t.extension)})
96 }
97 }
98 if len(out.Content) == 0 {
99 // Defensive: an empty content slice would confuse the model.
100 out.Content = []provider.Content{provider.TextBlock{Text: "(extension returned no content)"}}
101 }
102 out.Details = map[string]any{"extension": t.extension, "tool": t.name}
103 return out, nil
104}
105
106// decodeBase64 is a tiny wrapper around encoding/base64 so we can
107// validate the extension's image data in one place.
108func decodeBase64(s string) ([]byte, error) {
109 if s == "" {
110 return nil, nil
111 }
112 return base64DecodeStd(s)
113}