1package agent
2
3import (
4 "bufio"
5 "context"
6 "fmt"
7 "io"
8 "os"
9 "strings"
10 "time"
11
12 "github.com/patriceckhart/zot/internal/agent/extensions"
13 "github.com/patriceckhart/zot/internal/agent/modes"
14 "github.com/patriceckhart/zot/internal/auth"
15 "github.com/patriceckhart/zot/internal/core"
16 "github.com/patriceckhart/zot/internal/provider"
17 "github.com/patriceckhart/zot/internal/skills"
18 "github.com/patriceckhart/zot/internal/tui"
19)
20
21// interactiveExtHooks is a tiny adapter that lets the extension
22// manager call back into the Interactive instance built later in
23// runInteractive. The forward-declared *modes.Interactive is filled
24// in immediately after manager construction.
25type interactiveExtHooks struct {
26 ivPtr **modes.Interactive
27}
28
29func (h *interactiveExtHooks) iv() *modes.Interactive {
30 if h == nil || h.ivPtr == nil {
31 return nil
32 }
33 return *h.ivPtr
34}
35
36func (h *interactiveExtHooks) Notify(extName, level, message string) {
37 if iv := h.iv(); iv != nil {
38 iv.Notify(extName, level, message)
39 }
40}
41func (h *interactiveExtHooks) Submit(text string) {
42 if iv := h.iv(); iv != nil {
43 iv.Submit(text)
44 }
45}
46func (h *interactiveExtHooks) Insert(text string) {
47 if iv := h.iv(); iv != nil {
48 iv.Insert(text)
49 }
50}
51func (h *interactiveExtHooks) Display(extName, text string) {
52 if iv := h.iv(); iv != nil {
53 iv.Display(extName, text)
54 }
55}
56
57// extToolAdapter bridges *extensions.Manager to the
58// ExtensionToolSource interface declared in build.go (kept narrow to
59// avoid a build->extensions import cycle). One adapter instance per
60// run; used at every Resolve point so re-built agents pick up the
61// same set of extension tools.
62type extToolAdapter struct {
63 mgr *extensions.Manager
64}
65
66func (a *extToolAdapter) Tools() []ExtensionToolInfo {
67 infos := a.mgr.Tools()
68 out := make([]ExtensionToolInfo, len(infos))
69 for i, t := range infos {
70 out[i] = ExtensionToolInfo{
71 Extension: t.Extension,
72 Name: t.Name,
73 Description: t.Description,
74 Schema: t.Schema,
75 }
76 }
77 return out
78}
79
80func (a *extToolAdapter) NewExtensionTool(info ExtensionToolInfo) core.Tool {
81 return extensions.NewTool(a.mgr, extensions.ToolInfo{
82 Extension: info.Extension,
83 Name: info.Name,
84 Description: info.Description,
85 Schema: info.Schema,
86 })
87}
88
89// Run is the top-level entrypoint for the zot binary.
90func Run(rawArgs []string, version string) error {
91 // Subcommand router: `zot bot ...` is handled separately so the
92 // generic flag parser doesn't reject "bot" as a positional arg.
93 if handled, err := runBotCommand(rawArgs, version); handled {
94 return err
95 }
96 if handled, err := runExtCommand(rawArgs); handled {
97 return err
98 }
99 // `zot rpc` is shorthand for `zot --rpc` so third-party apps can
100 // spawn the binary with a clean argv. Strip the leading 'rpc'
101 // token and let the rest flow through the normal arg parser.
102 if len(rawArgs) > 0 && rawArgs[0] == "rpc" {
103 rawArgs = append([]string{"--rpc"}, rawArgs[1:]...)
104 }
105
106 args, err := ParseArgs(rawArgs)
107 if err != nil {
108 PrintHelp(version)
109 return err
110 }
111 if args.Help {
112 PrintHelp(version)
113 return nil
114 }
115 if args.Version {
116 fmt.Println("zot", version)
117 return nil
118 }
119 // Model catalog: load any cached discovery data before we inspect
120 // the model list (list-models, print/json, interactive).
121 LoadCachedModels()
122
123 if args.ListModels {
124 printModels()
125 return nil
126 }
127
128 ctx := context.Background()
129
130 // Kick an async refresh of the live model catalog. The first run of
131 // zot hits the network; subsequent runs within CacheTTL do nothing.
132 RefreshModelsAsync()
133
134 switch args.Mode {
135 case ModePrint:
136 return runPrintMode(ctx, args, version)
137 case ModeJSON:
138 return runJSONMode(ctx, args, version)
139 case ModeRPC:
140 return runRPCMode(ctx, args, version)
141 default:
142 return runInteractive(ctx, args, version)
143 }
144}
145
146// ---- print / json modes: require credentials, run single-shot ----
147
148func runPrintMode(ctx context.Context, args Args, version string) error {
149 r, err := Resolve(args, true)
150 if err != nil {
151 return err
152 }
153 ag := r.NewAgent()
154 sess, _ := openOrCreateSession(args, r, ag, version)
155 defer sess.Close()
156
157 prompt := args.Prompt
158 if prompt == "" {
159 piped, _ := readAllStdin()
160 prompt = strings.TrimSpace(piped)
161 }
162 if prompt == "" {
163 return fmt.Errorf("print mode requires a prompt (arg or stdin)")
164 }
165
166 start := len(ag.Messages())
167 err = modes.RunPrint(ctx, ag, prompt, nil, os.Stdout)
168 WriteNewTranscript(ag, sess, start)
169 return err
170}
171
172func runJSONMode(ctx context.Context, args Args, version string) error {
173 r, err := Resolve(args, true)
174 if err != nil {
175 return err
176 }
177 ag := r.NewAgent()
178 sess, _ := openOrCreateSession(args, r, ag, version)
179 defer sess.Close()
180
181 prompt := args.Prompt
182 if prompt == "" {
183 piped, _ := readAllStdin()
184 prompt = strings.TrimSpace(piped)
185 }
186 if prompt == "" {
187 return fmt.Errorf("json mode requires a prompt (arg or stdin)")
188 }
189
190 start := len(ag.Messages())
191 err = modes.RunJSON(ctx, ag, prompt, nil, os.Stdout)
192 WriteNewTranscript(ag, sess, start)
193 return err
194}
195
196// ---- interactive mode: opens the TUI even without credentials ----
197
198func runInteractive(ctx context.Context, args Args, version string) error {
199 // Resolve WITHOUT requiring credentials.
200 r, err := Resolve(args, false)
201 if err != nil {
202 return err
203 }
204
205 authStore := AuthStoreFor()
206 mgr := auth.NewManager(authStore)
207 defer mgr.Close()
208
209 // Keep the sandbox pointer stable across agent rebuilds (login / model
210 // switch). The Interactive UI toggles the lock via this pointer, and
211 // rebuilt tool instances must share the same one so the lock sticks.
212 sharedSandbox := r.Sandbox
213
214 // Build the extension manager BEFORE the agent so we can fold
215 // extension-defined tools into the registry. Forward-declare iv so
216 // the host hooks adapter can dereference it after construction.
217 var iv *modes.Interactive
218 extHooks := &interactiveExtHooks{ivPtr: &iv}
219 extMgr := extensions.New(ZotHome(), r.CWD, version, r.Provider, r.Model, extHooks)
220 discoveryErrs := extMgr.Discover(ctx)
221 for _, e := range discoveryErrs {
222 fmt.Fprintln(os.Stderr, "extension load:", e)
223 }
224 // Wait briefly for extensions to flush their initial register_tool
225 // frames before we build the agent's tool registry. Half a second
226 // is plenty for any extension that's actually well-behaved; ones
227 // that don't send a ready frame eat the full grace and proceed.
228 extMgr.WaitForReady(500 * time.Millisecond)
229 defer extMgr.Stop(2 * time.Second)
230
231 extToolAdapter := &extToolAdapter{mgr: extMgr}
232 r.MergeExtensionTools(extToolAdapter)
233
234 // Capture current args in a closure so BuildAgent can re-resolve
235 // after a successful login (picks up the newly stored credential).
236 buildAgent := func() (*core.Agent, string, string, error) {
237 resolved, err := Resolve(args, true)
238 if err != nil {
239 return nil, "", "", err
240 }
241 resolved.UseSandbox(sharedSandbox)
242 resolved.MergeExtensionTools(extToolAdapter)
243 return resolved.NewAgent(), resolved.Provider, resolved.Model, nil
244 }
245
246 // Rebuild agent with an explicit provider/model override.
247 buildAgentFor := func(providerOverride, modelOverride string) (*core.Agent, string, string, error) {
248 next := args
249 if providerOverride != "" {
250 next.Provider = providerOverride
251 }
252 if modelOverride != "" {
253 next.Model = modelOverride
254 }
255 resolved, err := Resolve(next, true)
256 if err != nil {
257 return nil, "", "", err
258 }
259 resolved.UseSandbox(sharedSandbox)
260 resolved.MergeExtensionTools(extToolAdapter)
261 return resolved.NewAgent(), resolved.Provider, resolved.Model, nil
262 }
263
264 var ag *core.Agent
265 if r.HasCredential() {
266 ag = r.NewAgent()
267 }
268
269 var sess *core.Session
270 var sessBaselineMsgs int // messages already on disk when current session opened
271 if !args.NoSess && ag != nil {
272 sess, _ = openOrCreateSession(args, r, ag, version)
273 if ag != nil {
274 sessBaselineMsgs = len(ag.Messages())
275 }
276 }
277 defer func() {
278 if sess != nil {
279 sess.Close()
280 }
281 }()
282
283 // loadSession replaces the current session with the one at path and
284 // hands its messages to the agent. Used by the /sessions picker.
285 loadSession := func(path string) error {
286 currentAg := ag // captured
287 if currentAg == nil {
288 return fmt.Errorf("no agent running; log in first")
289 }
290 newSess, msgs, err := core.OpenSession(path)
291 if err != nil {
292 return err
293 }
294 // Flush any unsaved messages to the old session before swapping.
295 if sess != nil {
296 WriteNewTranscript(currentAg, sess, sessBaselineMsgs)
297 _ = sess.Close()
298 }
299 sess = newSess
300 currentAg.SetMessages(msgs)
301 sessBaselineMsgs = len(msgs)
302 return nil
303 }
304
305 term := tui.NewProcTerm()
306
307 // Kick off the async update check so the banner can appear when the
308 // http response eventually arrives (usually <1s on cached DNS). Map
309 // agent.UpdateInfo -> modes.UpdateInfo here to avoid a cyclic import.
310 updateCh := make(chan modes.UpdateInfo, 1)
311 go func() {
312 defer close(updateCh)
313 src := <-CheckForUpdateAsync(ZotHome(), version)
314 updateCh <- modes.UpdateInfo{
315 Current: src.Current,
316 Latest: src.Latest,
317 Available: src.Available,
318 URL: src.URL,
319 }
320 }()
321
322 iv = modes.NewInteractive(modes.InteractiveConfig{
323 Terminal: term,
324 Theme: tui.Dark,
325 Model: r.Model,
326 Provider: r.Provider,
327 AuthMethod: r.AuthMethod,
328 BaseURL: r.BaseURL,
329 Reasoning: r.Reasoning,
330 SystemPrompt: r.SystemPrompt,
331 Tools: r.ToolRegistry,
332 MaxSteps: r.MaxSteps,
333 CWD: r.CWD,
334 ZotHome: ZotHome(),
335 Version: version,
336 UpdateInfoChan: updateCh,
337 Sandbox: sharedSandbox,
338 Agent: ag,
339 InitialInput: args.Prompt,
340 AuthManager: mgr,
341 BuildAgent: buildAgent,
342 BuildAgentFor: buildAgentFor,
343 LoadSession: loadSession,
344 Extensions: extMgr,
345 SkillSnapshot: func() []*skills.Skill {
346 // Re-discover so the picker reflects edits made during
347 // the session. Cheap; SKILL.md files are small.
348 userHome, _ := os.UserHomeDir()
349 list, _ := skills.Discover(ZotHome(), r.CWD, userHome)
350 return list
351 },
352 PersistModel: func(providerName, model string) {
353 // Update config.json so next launch uses the same pick.
354 cfg, _ := LoadConfig()
355 cfg.Provider = providerName
356 cfg.Model = model
357 _ = SaveConfig(cfg)
358 // Update the active session's meta so resume picks this up.
359 if sess != nil {
360 _ = sess.UpdateModel(providerName, model)
361 }
362 },
363 })
364
365 runErr := iv.Run(ctx)
366
367 // Flush final transcript to session (only if we had / ended up with an agent).
368 if finalAg := iv.Agent(); finalAg != nil && sess != nil {
369 WriteNewTranscript(finalAg, sess, sessBaselineMsgs)
370 }
371 return runErr
372}
373
374// openOrCreateSession returns a session for the run. sess may be nil
375// with a nil error if session persistence is disabled.
376func openOrCreateSession(args Args, r Resolved, ag *core.Agent, version string) (*core.Session, error) {
377 if args.NoSess {
378 return nil, nil
379 }
380 // Sweep meta-only files left over from older zot versions (and from
381 // any session that crashed before its first AppendMessage). Cheap;
382 // reads the first few bytes of each file in the cwd's session dir.
383 core.PruneEmptySessions(ZotHome(), args.CWD)
384 var (
385 s *core.Session
386 msgs []provider.Message
387 err error
388 )
389 switch {
390 case args.Session != "":
391 s, msgs, err = core.OpenSession(args.Session)
392 case args.Continue:
393 latest := core.LatestSession(ZotHome(), args.CWD)
394 if latest != "" {
395 s, msgs, err = core.OpenSession(latest)
396 }
397 case args.Resume:
398 picked, perr := pickSession(args.CWD)
399 if perr != nil {
400 return nil, perr
401 }
402 if picked != "" {
403 s, msgs, err = core.OpenSession(picked)
404 }
405 }
406 if err != nil {
407 return nil, err
408 }
409 if s != nil {
410 ag.SetMessages(msgs)
411 return s, nil
412 }
413 return core.NewSession(ZotHome(), args.CWD, r.Provider, r.Model, version)
414}
415
416func pickSession(cwd string) (string, error) {
417 files := core.ListSessions(ZotHome(), cwd)
418 if len(files) == 0 {
419 fmt.Fprintln(os.Stderr, "no sessions for", cwd)
420 return "", nil
421 }
422 for i, f := range files {
423 fmt.Fprintf(os.Stderr, " %2d) %s\n", i+1, f)
424 }
425 fmt.Fprint(os.Stderr, "pick #: ")
426 rd := bufio.NewReader(os.Stdin)
427 line, _ := rd.ReadString('\n')
428 line = strings.TrimSpace(line)
429 var n int
430 if _, err := fmt.Sscanf(line, "%d", &n); err != nil || n < 1 || n > len(files) {
431 return "", fmt.Errorf("invalid selection")
432 }
433 return files[n-1], nil
434}
435
436// WriteNewTranscript appends only messages after index `from` from the
437// agent's transcript to the session.
438func WriteNewTranscript(ag *core.Agent, sess *core.Session, from int) {
439 if sess == nil || ag == nil {
440 return
441 }
442 msgs := ag.Messages()
443 for i := from; i < len(msgs); i++ {
444 _ = sess.AppendMessage(msgs[i])
445 }
446 cum := ag.Cost()
447 _ = sess.AppendUsage(cum, cum)
448}
449
450func readAllStdin() (string, error) {
451 fi, err := os.Stdin.Stat()
452 if err != nil {
453 return "", err
454 }
455 if (fi.Mode() & os.ModeCharDevice) != 0 {
456 return "", nil
457 }
458 b, err := io.ReadAll(os.Stdin)
459 return string(b), err
460}
461
462func printModels() {
463 fmt.Println("provider model id context max-out reasoning source name")
464 for _, m := range provider.Active() {
465 reason := " "
466 if m.Reasoning {
467 reason = "✓"
468 }
469 source := m.Source
470 if source == "" {
471 source = "catalog"
472 }
473 if m.Speculative {
474 source = "speculative"
475 }
476 fmt.Printf("%-10s %-30s %8d %8d %s %-11s %s\n",
477 m.Provider, m.ID, m.ContextWindow, m.MaxOutput, reason, source, m.DisplayName)
478 }
479}