1package extensions
2
3import (
4 "bytes"
5 "context"
6 "encoding/json"
7 "slices"
8
9 "github.com/patriceckhart/zot/packages/agent/extproto"
10)
11
12// TrackToolCall records the initiator for confirmation and interception
13// events until the call finishes. The returned function removes the entry.
14func (m *Manager) TrackToolCall(id, origin string) func() {
15 m.mu.Lock()
16 if m.callOrigins == nil {
17 m.callOrigins = make(map[string]string)
18 }
19 m.callOrigins[id] = origin
20 m.mu.Unlock()
21 return func() {
22 m.mu.Lock()
23 delete(m.callOrigins, id)
24 m.mu.Unlock()
25 }
26}
27
28// ToolCallOrigin returns the initiating extension for an active host call.
29func (m *Manager) ToolCallOrigin(id string) string {
30 m.mu.RLock()
31 defer m.mu.RUnlock()
32 return m.callOrigins[id]
33}
34
35// The chain follows a call across extension processes through the host.
36// Only host-created parent IDs can carry it into a nested request.
37type toolChainKey struct{}
38
39type outboundCall struct {
40 ctx context.Context
41 chain []string
42}
43
44const maxToolCallDepth = 4
45
46// SetToolCaller binds host-managed execution to the current agent mode.
47// The callback must enforce the same policy as ordinary model tool calls.
48func (m *Manager) SetToolCaller(fn func(context.Context, string, string, json.RawMessage) extproto.ToolResultFromHost) {
49 m.mu.Lock()
50 m.callTool = fn
51 m.mu.Unlock()
52}
53
54func (m *Manager) handleCallTool(ctx context.Context, ext *Extension, req extproto.CallToolFromExt, cancel context.CancelFunc, parentFound, activeOutbound bool, chain []string) {
55 defer cancel()
56 defer func() {
57 ext.mu.Lock()
58 delete(ext.inboundCalls, req.ID)
59 ext.mu.Unlock()
60 }()
61 result := extproto.ToolResultFromHost{Type: "tool_result", ID: req.ID, Name: req.Name}
62 fail := func(message string) {
63 result.IsError = true
64 result.Content = []extproto.ContentBlock{{Type: "text", Text: message}}
65 }
66
67 m.mu.RLock()
68 caller := m.callTool
69 m.mu.RUnlock()
70 switch {
71 case !ext.callTool:
72 fail("extension did not opt in to call_tool")
73 case req.ParentID != "" && !parentFound:
74 fail("unknown or expired parent tool call")
75 case req.ParentID == "" && activeOutbound:
76 fail("parent_id is required while an extension tool is active")
77 case slices.Contains(chain, req.Name):
78 fail("tool call cycle detected")
79 case len(chain) >= maxToolCallDepth:
80 fail("tool call depth limit reached")
81 case caller == nil:
82 fail("host tool calls are unavailable")
83 case req.Name == "" || len(req.Args) == 0 || !json.Valid(req.Args) || bytes.TrimSpace(req.Args)[0] != '{':
84 fail("call_tool requires a name and JSON object args")
85 default:
86 ctx = context.WithValue(ctx, toolChainKey{}, append(append([]string(nil), chain...), req.Name))
87 result = caller(ctx, ext.Manifest.Name, req.Name, req.Args)
88 result.Type, result.ID, result.Name = "tool_result", req.ID, req.Name
89 }
90 frame, err := extproto.Encode(result)
91 if err != nil {
92 return
93 }
94 if pipe, ok := ext.stdin.(*orderedPipe); ok {
95 _ = pipe.enqueue(frame, nil)
96 } else if ext.stdin != nil {
97 _, _ = ext.stdin.Write(frame)
98 }
99}