1package agent
2
3import (
4 "bufio"
5 "context"
6 "encoding/json"
7 "fmt"
8 "io"
9 "os"
10 "strings"
11 "time"
12
13 "github.com/patriceckhart/zot/internal/agent/extensions"
14 "github.com/patriceckhart/zot/internal/agent/modes"
15 "github.com/patriceckhart/zot/internal/auth"
16 "github.com/patriceckhart/zot/internal/core"
17 "github.com/patriceckhart/zot/internal/extproto"
18 "github.com/patriceckhart/zot/internal/provider"
19 "github.com/patriceckhart/zot/internal/skills"
20 "github.com/patriceckhart/zot/internal/tui"
21)
22
23// interactiveExtHooks is a tiny adapter that lets the extension
24// manager call back into the Interactive instance built later in
25// runInteractive. The forward-declared *modes.Interactive is filled
26// in immediately after manager construction.
27type interactiveExtHooks struct {
28 ivPtr **modes.Interactive
29}
30
31func (h *interactiveExtHooks) iv() *modes.Interactive {
32 if h == nil || h.ivPtr == nil {
33 return nil
34 }
35 return *h.ivPtr
36}
37
38func (h *interactiveExtHooks) Notify(extName, level, message string) {
39 if iv := h.iv(); iv != nil {
40 iv.Notify(extName, level, message)
41 }
42}
43func (h *interactiveExtHooks) Submit(text string) {
44 if iv := h.iv(); iv != nil {
45 iv.Submit(text)
46 }
47}
48func (h *interactiveExtHooks) Insert(text string) {
49 if iv := h.iv(); iv != nil {
50 iv.Insert(text)
51 }
52}
53func (h *interactiveExtHooks) Display(extName, text string) {
54 if iv := h.iv(); iv != nil {
55 iv.Display(extName, text)
56 }
57}
58
59// extToolAdapter bridges *extensions.Manager to the
60// ExtensionToolSource interface declared in build.go (kept narrow to
61// avoid a build->extensions import cycle). One adapter instance per
62// run; used at every Resolve point so re-built agents pick up the
63// same set of extension tools.
64type extToolAdapter struct {
65 mgr *extensions.Manager
66}
67
68func (a *extToolAdapter) Tools() []ExtensionToolInfo {
69 infos := a.mgr.Tools()
70 out := make([]ExtensionToolInfo, len(infos))
71 for i, t := range infos {
72 out[i] = ExtensionToolInfo{
73 Extension: t.Extension,
74 Name: t.Name,
75 Description: t.Description,
76 Schema: t.Schema,
77 }
78 }
79 return out
80}
81
82func (a *extToolAdapter) NewExtensionTool(info ExtensionToolInfo) core.Tool {
83 return extensions.NewTool(a.mgr, extensions.ToolInfo{
84 Extension: info.Extension,
85 Name: info.Name,
86 Description: info.Description,
87 Schema: info.Schema,
88 })
89}
90
91// fanoutAgentEvent translates a core.AgentEvent into the wire-format
92// EventFromHost and pushes it through the extension manager. Only
93// the events that have a clear extension-facing meaning are
94// forwarded; internal-only ones (text_delta, tool_progress) are
95// dropped to keep the per-extension stream sane.
96func fanoutAgentEvent(mgr *extensions.Manager, ev core.AgentEvent) {
97 if mgr == nil {
98 return
99 }
100 switch e := ev.(type) {
101 case core.EvTurnStart:
102 mgr.EmitEvent(extproto.EventFromHost{Event: "turn_start", Step: e.Step})
103 case core.EvToolCall:
104 mgr.EmitEvent(extproto.EventFromHost{
105 Event: "tool_call", ToolID: e.ID, ToolName: e.Name, ToolArgs: e.Args,
106 })
107 case core.EvAssistantMessage:
108 // Concat the visible text portions of the message; binary
109 // blocks (tool_use, etc.) are skipped because subscribers
110 // usually want a string they can grep / display.
111 var text string
112 for _, c := range e.Message.Content {
113 if tb, ok := c.(provider.TextBlock); ok {
114 text += tb.Text
115 }
116 }
117 mgr.EmitEvent(extproto.EventFromHost{Event: "assistant_message", Text: text})
118 case core.EvTurnEnd:
119 ev := extproto.EventFromHost{Event: "turn_end", Stop: string(e.Stop)}
120 if e.Err != nil {
121 ev.Error = e.Err.Error()
122 }
123 mgr.EmitEvent(ev)
124 }
125}
126
127// Run is the top-level entrypoint for the zot binary.
128func Run(rawArgs []string, version string) error {
129 // Subcommand router: `zot bot ...` is handled separately so the
130 // generic flag parser doesn't reject "bot" as a positional arg.
131 if handled, err := runBotCommand(rawArgs, version); handled {
132 return err
133 }
134 if handled, err := runExtCommand(rawArgs); handled {
135 return err
136 }
137 // `zot rpc` is shorthand for `zot --rpc` so third-party apps can
138 // spawn the binary with a clean argv. Strip the leading 'rpc'
139 // token and let the rest flow through the normal arg parser.
140 if len(rawArgs) > 0 && rawArgs[0] == "rpc" {
141 rawArgs = append([]string{"--rpc"}, rawArgs[1:]...)
142 }
143
144 args, err := ParseArgs(rawArgs)
145 if err != nil {
146 PrintHelp(version)
147 return err
148 }
149 if args.Help {
150 PrintHelp(version)
151 return nil
152 }
153 if args.Version {
154 fmt.Println("zot", version)
155 return nil
156 }
157 // Model catalog: load any cached discovery data before we inspect
158 // the model list (list-models, print/json, interactive).
159 LoadCachedModels()
160
161 if args.ListModels {
162 printModels()
163 return nil
164 }
165
166 ctx := context.Background()
167
168 // Kick an async refresh of the live model catalog. The first run of
169 // zot hits the network; subsequent runs within CacheTTL do nothing.
170 RefreshModelsAsync()
171
172 switch args.Mode {
173 case ModePrint:
174 return runPrintMode(ctx, args, version)
175 case ModeJSON:
176 return runJSONMode(ctx, args, version)
177 case ModeRPC:
178 return runRPCMode(ctx, args, version)
179 default:
180 return runInteractive(ctx, args, version)
181 }
182}
183
184// ---- print / json modes: require credentials, run single-shot ----
185
186func runPrintMode(ctx context.Context, args Args, version string) error {
187 r, err := Resolve(args, true)
188 if err != nil {
189 return err
190 }
191 ag := r.NewAgent()
192 wireNoYoloAutoRefuse(ag, args)
193 sess, _ := openOrCreateSession(args, r, ag, version)
194 defer sess.Close()
195
196 prompt := args.Prompt
197 if prompt == "" {
198 piped, _ := readAllStdin()
199 prompt = strings.TrimSpace(piped)
200 }
201 if prompt == "" {
202 return fmt.Errorf("print mode requires a prompt (arg or stdin)")
203 }
204
205 start := len(ag.Messages())
206 err = modes.RunPrint(ctx, ag, prompt, nil, os.Stdout)
207 WriteNewTranscript(ag, sess, start)
208 return err
209}
210
211// wireNoYoloAutoRefuse installs a BeforeToolExecute hook that
212// refuses every tool call when --no-yolo is active but there's no
213// interactive UI to prompt (print / json / rpc modes). The reason
214// is written in a way the model can learn from so it proposes a
215// different action rather than looping on the same tool.
216func wireNoYoloAutoRefuse(ag *core.Agent, args Args) {
217 if !args.NoYolo || ag == nil {
218 return
219 }
220 gate := core.NewConfirmGate(nil) // nil inner = auto-refuse
221 prev := ag.BeforeToolExecute
222 ag.BeforeToolExecute = func(call provider.ToolCallBlock) (bool, string, json.RawMessage) {
223 ok, reason, _ := gate.Check(call.Name, core.BuildPreview(call.Arguments, 120))
224 if !ok {
225 return false, reason, nil
226 }
227 if prev != nil {
228 return prev(call)
229 }
230 return true, "", nil
231 }
232}
233
234func runJSONMode(ctx context.Context, args Args, version string) error {
235 r, err := Resolve(args, true)
236 if err != nil {
237 return err
238 }
239 ag := r.NewAgent()
240 wireNoYoloAutoRefuse(ag, args)
241 sess, _ := openOrCreateSession(args, r, ag, version)
242 defer sess.Close()
243
244 prompt := args.Prompt
245 if prompt == "" {
246 piped, _ := readAllStdin()
247 prompt = strings.TrimSpace(piped)
248 }
249 if prompt == "" {
250 return fmt.Errorf("json mode requires a prompt (arg or stdin)")
251 }
252
253 start := len(ag.Messages())
254 err = modes.RunJSON(ctx, ag, prompt, nil, os.Stdout)
255 WriteNewTranscript(ag, sess, start)
256 return err
257}
258
259// ---- interactive mode: opens the TUI even without credentials ----
260
261func runInteractive(ctx context.Context, args Args, version string) error {
262 // Resolve WITHOUT requiring credentials.
263 r, err := Resolve(args, false)
264 if err != nil {
265 return err
266 }
267
268 authStore := AuthStoreFor()
269 mgr := auth.NewManager(authStore)
270 defer mgr.Close()
271
272 // Keep the sandbox pointer stable across agent rebuilds (login / model
273 // switch). The Interactive UI toggles the lock via this pointer, and
274 // rebuilt tool instances must share the same one so the lock sticks.
275 sharedSandbox := r.Sandbox
276
277 // Build the extension manager BEFORE the agent so we can fold
278 // extension-defined tools into the registry. Forward-declare iv so
279 // the host hooks adapter can dereference it after construction.
280 var iv *modes.Interactive
281 extHooks := &interactiveExtHooks{ivPtr: &iv}
282 extMgr := extensions.New(ZotHome(), r.CWD, version, r.Provider, r.Model, extHooks)
283 // --ext paths first so they win against installed extensions of
284 // the same name (loadOne's first-write-wins semantics).
285 for _, e := range extMgr.LoadExplicit(ctx, args.Exts) {
286 fmt.Fprintln(os.Stderr, "extension load:", e)
287 }
288 // --no-ext skips the global + project-local discovery scan;
289 // explicit --ext paths above are still honoured so you can run
290 // "only this extension" with --no-ext --ext ./x.
291 if !args.NoExt {
292 for _, e := range extMgr.Discover(ctx) {
293 fmt.Fprintln(os.Stderr, "extension load:", e)
294 }
295 }
296 // Wait briefly for extensions to flush their initial register_tool
297 // frames before we build the agent's tool registry. Half a second
298 // is plenty for any extension that's actually well-behaved; ones
299 // that don't send a ready frame eat the full grace and proceed.
300 // 3s is the per-extension grace period for the ready frame.
301 // Native binaries are instant; runtimes like `npx tsx` take ~1.5s
302 // from cold cache. The wait is tight only for extensions that
303 // haven't sent ready by then; ones that signalled earlier release
304 // the wait immediately.
305 extMgr.WaitForReady(3 * time.Second)
306 defer extMgr.Stop(2 * time.Second)
307
308 extToolAdapter := &extToolAdapter{mgr: extMgr}
309 r.MergeExtensionTools(extToolAdapter)
310
311 // Confirmation gate: when --no-yolo is on, the agent must ask
312 // the user before every tool call. In interactive mode the TUI
313 // provides the Confirmer; in print/json/rpc modes there's no
314 // way to prompt, so the gate is constructed with a nil inner
315 // which auto-refuses every call with a helpful reason.
316 var confirmGate *core.ConfirmGate
317 if args.NoYolo {
318 confirmGate = core.NewConfirmGate(nil) // set below for interactive
319 }
320
321 // Capture current args in a closure so BuildAgent can re-resolve
322 // after a successful login (picks up the newly stored credential).
323 wireAgentExt := func(a *core.Agent) *core.Agent {
324 if a == nil {
325 return a
326 }
327 a.BeforeToolExecute = func(call provider.ToolCallBlock) (bool, string, json.RawMessage) {
328 // Confirm gate runs FIRST: if the user refused, we don't
329 // waste extension-intercept time or let guards see the call.
330 if confirmGate != nil {
331 ok, reason, _ := confirmGate.Check(call.Name, core.BuildPreview(call.Arguments, 120))
332 if !ok {
333 return false, reason, nil
334 }
335 }
336 r := extMgr.InterceptToolCall(ctx, call.ID, call.Name, call.Arguments)
337 if r.Block {
338 return false, r.Reason, nil
339 }
340 return true, "", r.ModifiedArgs
341 }
342 a.BeforeTurn = func(step int) (bool, string) {
343 r := extMgr.InterceptTurnStart(ctx, step)
344 return !r.Block, r.Reason
345 }
346 a.BeforeAssistantMessage = func(text string) (bool, string, string) {
347 r := extMgr.InterceptAssistantMessage(ctx, text)
348 if r.Block {
349 return false, r.Reason, ""
350 }
351 return true, "", r.ReplaceText
352 }
353 a.OnEvent = func(ev core.AgentEvent) { fanoutAgentEvent(extMgr, ev) }
354 return a
355 }
356
357 buildAgent := func() (*core.Agent, string, string, error) {
358 resolved, err := Resolve(args, true)
359 if err != nil {
360 return nil, "", "", err
361 }
362 resolved.UseSandbox(sharedSandbox)
363 resolved.MergeExtensionTools(extToolAdapter)
364 return wireAgentExt(resolved.NewAgent()), resolved.Provider, resolved.Model, nil
365 }
366
367 // Rebuild agent with an explicit provider/model override.
368 buildAgentFor := func(providerOverride, modelOverride string) (*core.Agent, string, string, error) {
369 next := args
370 if providerOverride != "" {
371 next.Provider = providerOverride
372 }
373 if modelOverride != "" {
374 next.Model = modelOverride
375 }
376 resolved, err := Resolve(next, true)
377 if err != nil {
378 return nil, "", "", err
379 }
380 resolved.UseSandbox(sharedSandbox)
381 resolved.MergeExtensionTools(extToolAdapter)
382 return wireAgentExt(resolved.NewAgent()), resolved.Provider, resolved.Model, nil
383 }
384
385 var ag *core.Agent
386 if r.HasCredential() {
387 ag = wireAgentExt(r.NewAgent())
388 }
389
390 // /reload-ext callback: after the manager has respawned every
391 // extension, re-resolve the tool registry (built-ins + freshly-
392 // registered extension tools) and swap it onto the current
393 // agent in-place. The current agent may have been replaced by a
394 // /model swap since spawn, so re-read the live `ag` on each
395 // invocation.
396 extMgr.SetOnReload(func() {
397 current := ag
398 if current == nil {
399 return
400 }
401 resolved, err := Resolve(args, true)
402 if err != nil {
403 return
404 }
405 resolved.UseSandbox(sharedSandbox)
406 resolved.MergeExtensionTools(extToolAdapter)
407 current.SetTools(resolved.ToolRegistry)
408 })
409
410 // Fire session_start once we know the manager's running.
411 extMgr.EmitEvent(extproto.EventFromHost{Event: "session_start"})
412
413 var sess *core.Session
414 var sessBaselineMsgs int // messages already on disk when current session opened
415 if !args.NoSess && ag != nil {
416 sess, _ = openOrCreateSession(args, r, ag, version)
417 if ag != nil {
418 sessBaselineMsgs = len(ag.Messages())
419 }
420 }
421 defer func() {
422 if sess != nil {
423 sess.Close()
424 }
425 }()
426
427 // loadSession replaces the current session with the one at path and
428 // hands its messages to the agent. Used by the /sessions picker.
429 loadSession := func(path string) error {
430 currentAg := ag // captured
431 if currentAg == nil {
432 return fmt.Errorf("no agent running; log in first")
433 }
434 newSess, msgs, err := core.OpenSession(path)
435 if err != nil {
436 return err
437 }
438 // Flush any unsaved messages to the old session before swapping.
439 if sess != nil {
440 WriteNewTranscript(currentAg, sess, sessBaselineMsgs)
441 _ = sess.Close()
442 }
443 sess = newSess
444 currentAg.SetMessages(msgs)
445 sessBaselineMsgs = len(msgs)
446 return nil
447 }
448
449 term := tui.NewProcTerm()
450
451 // Kick off the async update check so the banner can appear when the
452 // http response eventually arrives (usually <1s on cached DNS). Map
453 // agent.UpdateInfo -> modes.UpdateInfo here to avoid a cyclic import.
454 updateCh := make(chan modes.UpdateInfo, 1)
455 go func() {
456 defer close(updateCh)
457 src := <-CheckForUpdateAsync(ZotHome(), version)
458 updateCh <- modes.UpdateInfo{
459 Current: src.Current,
460 Latest: src.Latest,
461 Available: src.Available,
462 URL: src.URL,
463 }
464 }()
465
466 // Changelog: when the running version differs from the last
467 // version whose release notes the user dismissed, fetch the
468 // release body from GitHub and have the TUI show it once. On
469 // first-ever launch (no prior LastChangelogShown), seed the
470 // stored version silently — don't dump release notes at someone
471 // who just installed.
472 changelogCh := make(chan modes.ChangelogPayload, 1)
473 go func() {
474 defer close(changelogCh)
475 cfg, _ := LoadConfig()
476 if cfg.LastChangelogShown == "" {
477 SeedChangelogVersion(version)
478 return
479 }
480 if !ShouldShowChangelog(version, cfg) {
481 return
482 }
483 info := <-FetchChangelogAsync(version)
484 if info.Body == "" {
485 return
486 }
487 changelogCh <- modes.ChangelogPayload{
488 Version: info.Version,
489 Body: info.Body,
490 URL: info.URL,
491 }
492 }()
493
494 iv = modes.NewInteractive(modes.InteractiveConfig{
495 Terminal: term,
496 Theme: tui.Dark,
497 Model: r.Model,
498 Provider: r.Provider,
499 AuthMethod: r.AuthMethod,
500 BaseURL: r.BaseURL,
501 Reasoning: r.Reasoning,
502 SystemPrompt: r.SystemPrompt,
503 Tools: r.ToolRegistry,
504 MaxSteps: r.MaxSteps,
505 CWD: r.CWD,
506 ZotHome: ZotHome(),
507 Version: version,
508 UpdateInfoChan: updateCh,
509 Sandbox: sharedSandbox,
510 Agent: ag,
511 InitialInput: args.Prompt,
512 AuthManager: mgr,
513 BuildAgent: buildAgent,
514 BuildAgentFor: buildAgentFor,
515 LoadSession: loadSession,
516 Extensions: extMgr,
517 ChangelogChan: changelogCh,
518 OnChangelogDismiss: func() {
519 _ = MarkChangelogShown(version)
520 },
521 SkillSnapshot: func() []*skills.Skill {
522 if args.NoSkill {
523 // --no-skill: nothing for the picker to show.
524 return nil
525 }
526 // Re-discover so the picker reflects edits made during
527 // the session. Cheap; SKILL.md files are small. Filter
528 // out built-in skills — they're hidden from user-facing
529 // surfaces because they're implementation detail; the
530 // model still sees them through the system-prompt
531 // manifest and the skill tool. User skills only appear
532 // when --with-skills is set; without it the picker shows
533 // nothing but the model still has the built-ins.
534 userHome, _ := os.UserHomeDir()
535 list, _ := skills.Discover(ZotHome(), r.CWD, userHome, args.WithSkills)
536 return skills.VisibleSkills(list)
537 },
538 NoYolo: args.NoYolo,
539 ConfirmGate: confirmGate,
540 PersistModel: func(providerName, model string) {
541 // Update config.json so next launch uses the same pick.
542 cfg, _ := LoadConfig()
543 cfg.Provider = providerName
544 cfg.Model = model
545 _ = SaveConfig(cfg)
546 // Update the active session's meta so resume picks this up.
547 if sess != nil {
548 _ = sess.UpdateModel(providerName, model)
549 }
550 },
551 })
552
553 // Bind the interactive TUI as the Confirmer. We deferred this
554 // until now because the gate is constructed before the TUI
555 // (the BeforeToolExecute closure captures it). SetConfirmer
556 // is mutex-guarded on the gate so this is safe.
557 if confirmGate != nil {
558 confirmGate.SetConfirmer(iv)
559 }
560
561 runErr := iv.Run(ctx)
562
563 // Flush final transcript to session (only if we had / ended up with an agent).
564 if finalAg := iv.Agent(); finalAg != nil && sess != nil {
565 WriteNewTranscript(finalAg, sess, sessBaselineMsgs)
566 }
567 return runErr
568}
569
570// openOrCreateSession returns a session for the run. sess may be nil
571// with a nil error if session persistence is disabled.
572func openOrCreateSession(args Args, r Resolved, ag *core.Agent, version string) (*core.Session, error) {
573 if args.NoSess {
574 return nil, nil
575 }
576 // Sweep meta-only files left over from older zot versions (and from
577 // any session that crashed before its first AppendMessage). Cheap;
578 // reads the first few bytes of each file in the cwd's session dir.
579 core.PruneEmptySessions(ZotHome(), args.CWD)
580 var (
581 s *core.Session
582 msgs []provider.Message
583 err error
584 )
585 switch {
586 case args.Session != "":
587 s, msgs, err = core.OpenSession(args.Session)
588 case args.Continue:
589 latest := core.LatestSession(ZotHome(), args.CWD)
590 if latest != "" {
591 s, msgs, err = core.OpenSession(latest)
592 }
593 case args.Resume:
594 picked, perr := pickSession(args.CWD)
595 if perr != nil {
596 return nil, perr
597 }
598 if picked != "" {
599 s, msgs, err = core.OpenSession(picked)
600 }
601 }
602 if err != nil {
603 return nil, err
604 }
605 if s != nil {
606 ag.SetMessages(msgs)
607 return s, nil
608 }
609 return core.NewSession(ZotHome(), args.CWD, r.Provider, r.Model, version)
610}
611
612func pickSession(cwd string) (string, error) {
613 files := core.ListSessions(ZotHome(), cwd)
614 if len(files) == 0 {
615 fmt.Fprintln(os.Stderr, "no sessions for", cwd)
616 return "", nil
617 }
618 for i, f := range files {
619 fmt.Fprintf(os.Stderr, " %2d) %s\n", i+1, f)
620 }
621 fmt.Fprint(os.Stderr, "pick #: ")
622 rd := bufio.NewReader(os.Stdin)
623 line, _ := rd.ReadString('\n')
624 line = strings.TrimSpace(line)
625 var n int
626 if _, err := fmt.Sscanf(line, "%d", &n); err != nil || n < 1 || n > len(files) {
627 return "", fmt.Errorf("invalid selection")
628 }
629 return files[n-1], nil
630}
631
632// WriteNewTranscript appends only messages after index `from` from the
633// agent's transcript to the session.
634func WriteNewTranscript(ag *core.Agent, sess *core.Session, from int) {
635 if sess == nil || ag == nil {
636 return
637 }
638 msgs := ag.Messages()
639 for i := from; i < len(msgs); i++ {
640 _ = sess.AppendMessage(msgs[i])
641 }
642 cum := ag.Cost()
643 _ = sess.AppendUsage(cum, cum)
644}
645
646func readAllStdin() (string, error) {
647 fi, err := os.Stdin.Stat()
648 if err != nil {
649 return "", err
650 }
651 if (fi.Mode() & os.ModeCharDevice) != 0 {
652 return "", nil
653 }
654 b, err := io.ReadAll(os.Stdin)
655 return string(b), err
656}
657
658func printModels() {
659 fmt.Println("provider model id context max-out reasoning source name")
660 for _, m := range provider.Active() {
661 reason := " "
662 if m.Reasoning {
663 reason = "✓"
664 }
665 source := m.Source
666 if source == "" {
667 source = "catalog"
668 }
669 if m.Speculative {
670 source = "speculative"
671 }
672 fmt.Printf("%-10s %-30s %8d %8d %s %-11s %s\n",
673 m.Provider, m.ID, m.ContextWindow, m.MaxOutput, reason, source, m.DisplayName)
674 }
675}