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 sess, _ := openOrCreateSession(args, r, ag, version)
193 defer sess.Close()
194
195 prompt := args.Prompt
196 if prompt == "" {
197 piped, _ := readAllStdin()
198 prompt = strings.TrimSpace(piped)
199 }
200 if prompt == "" {
201 return fmt.Errorf("print mode requires a prompt (arg or stdin)")
202 }
203
204 start := len(ag.Messages())
205 err = modes.RunPrint(ctx, ag, prompt, nil, os.Stdout)
206 WriteNewTranscript(ag, sess, start)
207 return err
208}
209
210func runJSONMode(ctx context.Context, args Args, version string) error {
211 r, err := Resolve(args, true)
212 if err != nil {
213 return err
214 }
215 ag := r.NewAgent()
216 sess, _ := openOrCreateSession(args, r, ag, version)
217 defer sess.Close()
218
219 prompt := args.Prompt
220 if prompt == "" {
221 piped, _ := readAllStdin()
222 prompt = strings.TrimSpace(piped)
223 }
224 if prompt == "" {
225 return fmt.Errorf("json mode requires a prompt (arg or stdin)")
226 }
227
228 start := len(ag.Messages())
229 err = modes.RunJSON(ctx, ag, prompt, nil, os.Stdout)
230 WriteNewTranscript(ag, sess, start)
231 return err
232}
233
234// ---- interactive mode: opens the TUI even without credentials ----
235
236func runInteractive(ctx context.Context, args Args, version string) error {
237 // Resolve WITHOUT requiring credentials.
238 r, err := Resolve(args, false)
239 if err != nil {
240 return err
241 }
242
243 authStore := AuthStoreFor()
244 mgr := auth.NewManager(authStore)
245 defer mgr.Close()
246
247 // Keep the sandbox pointer stable across agent rebuilds (login / model
248 // switch). The Interactive UI toggles the lock via this pointer, and
249 // rebuilt tool instances must share the same one so the lock sticks.
250 sharedSandbox := r.Sandbox
251
252 // Build the extension manager BEFORE the agent so we can fold
253 // extension-defined tools into the registry. Forward-declare iv so
254 // the host hooks adapter can dereference it after construction.
255 var iv *modes.Interactive
256 extHooks := &interactiveExtHooks{ivPtr: &iv}
257 extMgr := extensions.New(ZotHome(), r.CWD, version, r.Provider, r.Model, extHooks)
258 // --ext paths first so they win against installed extensions of
259 // the same name (loadOne's first-write-wins semantics).
260 for _, e := range extMgr.LoadExplicit(ctx, args.Exts) {
261 fmt.Fprintln(os.Stderr, "extension load:", e)
262 }
263 // --no-ext skips the global + project-local discovery scan;
264 // explicit --ext paths above are still honoured so you can run
265 // "only this extension" with --no-ext --ext ./x.
266 if !args.NoExt {
267 for _, e := range extMgr.Discover(ctx) {
268 fmt.Fprintln(os.Stderr, "extension load:", e)
269 }
270 }
271 // Wait briefly for extensions to flush their initial register_tool
272 // frames before we build the agent's tool registry. Half a second
273 // is plenty for any extension that's actually well-behaved; ones
274 // that don't send a ready frame eat the full grace and proceed.
275 // 3s is the per-extension grace period for the ready frame.
276 // Native binaries are instant; runtimes like `npx tsx` take ~1.5s
277 // from cold cache. The wait is tight only for extensions that
278 // haven't sent ready by then; ones that signalled earlier release
279 // the wait immediately.
280 extMgr.WaitForReady(3 * time.Second)
281 defer extMgr.Stop(2 * time.Second)
282
283 extToolAdapter := &extToolAdapter{mgr: extMgr}
284 r.MergeExtensionTools(extToolAdapter)
285
286 // Capture current args in a closure so BuildAgent can re-resolve
287 // after a successful login (picks up the newly stored credential).
288 wireAgentExt := func(a *core.Agent) *core.Agent {
289 if a == nil {
290 return a
291 }
292 a.BeforeToolExecute = func(call provider.ToolCallBlock) (bool, string, json.RawMessage) {
293 r := extMgr.InterceptToolCall(ctx, call.ID, call.Name, call.Arguments)
294 if r.Block {
295 return false, r.Reason, nil
296 }
297 return true, "", r.ModifiedArgs
298 }
299 a.BeforeTurn = func(step int) (bool, string) {
300 r := extMgr.InterceptTurnStart(ctx, step)
301 return !r.Block, r.Reason
302 }
303 a.BeforeAssistantMessage = func(text string) (bool, string, string) {
304 r := extMgr.InterceptAssistantMessage(ctx, text)
305 if r.Block {
306 return false, r.Reason, ""
307 }
308 return true, "", r.ReplaceText
309 }
310 a.OnEvent = func(ev core.AgentEvent) { fanoutAgentEvent(extMgr, ev) }
311 return a
312 }
313
314 buildAgent := func() (*core.Agent, string, string, error) {
315 resolved, err := Resolve(args, true)
316 if err != nil {
317 return nil, "", "", err
318 }
319 resolved.UseSandbox(sharedSandbox)
320 resolved.MergeExtensionTools(extToolAdapter)
321 return wireAgentExt(resolved.NewAgent()), resolved.Provider, resolved.Model, nil
322 }
323
324 // Rebuild agent with an explicit provider/model override.
325 buildAgentFor := func(providerOverride, modelOverride string) (*core.Agent, string, string, error) {
326 next := args
327 if providerOverride != "" {
328 next.Provider = providerOverride
329 }
330 if modelOverride != "" {
331 next.Model = modelOverride
332 }
333 resolved, err := Resolve(next, true)
334 if err != nil {
335 return nil, "", "", err
336 }
337 resolved.UseSandbox(sharedSandbox)
338 resolved.MergeExtensionTools(extToolAdapter)
339 return wireAgentExt(resolved.NewAgent()), resolved.Provider, resolved.Model, nil
340 }
341
342 var ag *core.Agent
343 if r.HasCredential() {
344 ag = wireAgentExt(r.NewAgent())
345 }
346
347 // /reload-ext callback: after the manager has respawned every
348 // extension, re-resolve the tool registry (built-ins + freshly-
349 // registered extension tools) and swap it onto the current
350 // agent in-place. The current agent may have been replaced by a
351 // /model swap since spawn, so re-read the live `ag` on each
352 // invocation.
353 extMgr.SetOnReload(func() {
354 current := ag
355 if current == nil {
356 return
357 }
358 resolved, err := Resolve(args, true)
359 if err != nil {
360 return
361 }
362 resolved.UseSandbox(sharedSandbox)
363 resolved.MergeExtensionTools(extToolAdapter)
364 current.SetTools(resolved.ToolRegistry)
365 })
366
367 // Fire session_start once we know the manager's running.
368 extMgr.EmitEvent(extproto.EventFromHost{Event: "session_start"})
369
370 var sess *core.Session
371 var sessBaselineMsgs int // messages already on disk when current session opened
372 if !args.NoSess && ag != nil {
373 sess, _ = openOrCreateSession(args, r, ag, version)
374 if ag != nil {
375 sessBaselineMsgs = len(ag.Messages())
376 }
377 }
378 defer func() {
379 if sess != nil {
380 sess.Close()
381 }
382 }()
383
384 // loadSession replaces the current session with the one at path and
385 // hands its messages to the agent. Used by the /sessions picker.
386 loadSession := func(path string) error {
387 currentAg := ag // captured
388 if currentAg == nil {
389 return fmt.Errorf("no agent running; log in first")
390 }
391 newSess, msgs, err := core.OpenSession(path)
392 if err != nil {
393 return err
394 }
395 // Flush any unsaved messages to the old session before swapping.
396 if sess != nil {
397 WriteNewTranscript(currentAg, sess, sessBaselineMsgs)
398 _ = sess.Close()
399 }
400 sess = newSess
401 currentAg.SetMessages(msgs)
402 sessBaselineMsgs = len(msgs)
403 return nil
404 }
405
406 term := tui.NewProcTerm()
407
408 // Kick off the async update check so the banner can appear when the
409 // http response eventually arrives (usually <1s on cached DNS). Map
410 // agent.UpdateInfo -> modes.UpdateInfo here to avoid a cyclic import.
411 updateCh := make(chan modes.UpdateInfo, 1)
412 go func() {
413 defer close(updateCh)
414 src := <-CheckForUpdateAsync(ZotHome(), version)
415 updateCh <- modes.UpdateInfo{
416 Current: src.Current,
417 Latest: src.Latest,
418 Available: src.Available,
419 URL: src.URL,
420 }
421 }()
422
423 // Changelog: when the running version differs from the last
424 // version whose release notes the user dismissed, fetch the
425 // release body from GitHub and have the TUI show it once. On
426 // first-ever launch (no prior LastChangelogShown), seed the
427 // stored version silently — don't dump release notes at someone
428 // who just installed.
429 changelogCh := make(chan modes.ChangelogPayload, 1)
430 go func() {
431 defer close(changelogCh)
432 cfg, _ := LoadConfig()
433 if cfg.LastChangelogShown == "" {
434 SeedChangelogVersion(version)
435 return
436 }
437 if !ShouldShowChangelog(version, cfg) {
438 return
439 }
440 info := <-FetchChangelogAsync(version)
441 if info.Body == "" {
442 return
443 }
444 changelogCh <- modes.ChangelogPayload{
445 Version: info.Version,
446 Body: info.Body,
447 URL: info.URL,
448 }
449 }()
450
451 iv = modes.NewInteractive(modes.InteractiveConfig{
452 Terminal: term,
453 Theme: tui.Dark,
454 Model: r.Model,
455 Provider: r.Provider,
456 AuthMethod: r.AuthMethod,
457 BaseURL: r.BaseURL,
458 Reasoning: r.Reasoning,
459 SystemPrompt: r.SystemPrompt,
460 Tools: r.ToolRegistry,
461 MaxSteps: r.MaxSteps,
462 CWD: r.CWD,
463 ZotHome: ZotHome(),
464 Version: version,
465 UpdateInfoChan: updateCh,
466 Sandbox: sharedSandbox,
467 Agent: ag,
468 InitialInput: args.Prompt,
469 AuthManager: mgr,
470 BuildAgent: buildAgent,
471 BuildAgentFor: buildAgentFor,
472 LoadSession: loadSession,
473 Extensions: extMgr,
474 ChangelogChan: changelogCh,
475 OnChangelogDismiss: func() {
476 _ = MarkChangelogShown(version)
477 },
478 SkillSnapshot: func() []*skills.Skill {
479 if args.NoSkill {
480 // --no-skill: nothing for the picker to show.
481 return nil
482 }
483 // Re-discover so the picker reflects edits made during
484 // the session. Cheap; SKILL.md files are small. Filter
485 // out built-in skills — they're hidden from user-facing
486 // surfaces because they're implementation detail; the
487 // model still sees them through the system-prompt
488 // manifest and the skill tool. User skills only appear
489 // when --with-skills is set; without it the picker shows
490 // nothing but the model still has the built-ins.
491 userHome, _ := os.UserHomeDir()
492 list, _ := skills.Discover(ZotHome(), r.CWD, userHome, args.WithSkills)
493 return skills.VisibleSkills(list)
494 },
495 PersistModel: func(providerName, model string) {
496 // Update config.json so next launch uses the same pick.
497 cfg, _ := LoadConfig()
498 cfg.Provider = providerName
499 cfg.Model = model
500 _ = SaveConfig(cfg)
501 // Update the active session's meta so resume picks this up.
502 if sess != nil {
503 _ = sess.UpdateModel(providerName, model)
504 }
505 },
506 })
507
508 runErr := iv.Run(ctx)
509
510 // Flush final transcript to session (only if we had / ended up with an agent).
511 if finalAg := iv.Agent(); finalAg != nil && sess != nil {
512 WriteNewTranscript(finalAg, sess, sessBaselineMsgs)
513 }
514 return runErr
515}
516
517// openOrCreateSession returns a session for the run. sess may be nil
518// with a nil error if session persistence is disabled.
519func openOrCreateSession(args Args, r Resolved, ag *core.Agent, version string) (*core.Session, error) {
520 if args.NoSess {
521 return nil, nil
522 }
523 // Sweep meta-only files left over from older zot versions (and from
524 // any session that crashed before its first AppendMessage). Cheap;
525 // reads the first few bytes of each file in the cwd's session dir.
526 core.PruneEmptySessions(ZotHome(), args.CWD)
527 var (
528 s *core.Session
529 msgs []provider.Message
530 err error
531 )
532 switch {
533 case args.Session != "":
534 s, msgs, err = core.OpenSession(args.Session)
535 case args.Continue:
536 latest := core.LatestSession(ZotHome(), args.CWD)
537 if latest != "" {
538 s, msgs, err = core.OpenSession(latest)
539 }
540 case args.Resume:
541 picked, perr := pickSession(args.CWD)
542 if perr != nil {
543 return nil, perr
544 }
545 if picked != "" {
546 s, msgs, err = core.OpenSession(picked)
547 }
548 }
549 if err != nil {
550 return nil, err
551 }
552 if s != nil {
553 ag.SetMessages(msgs)
554 return s, nil
555 }
556 return core.NewSession(ZotHome(), args.CWD, r.Provider, r.Model, version)
557}
558
559func pickSession(cwd string) (string, error) {
560 files := core.ListSessions(ZotHome(), cwd)
561 if len(files) == 0 {
562 fmt.Fprintln(os.Stderr, "no sessions for", cwd)
563 return "", nil
564 }
565 for i, f := range files {
566 fmt.Fprintf(os.Stderr, " %2d) %s\n", i+1, f)
567 }
568 fmt.Fprint(os.Stderr, "pick #: ")
569 rd := bufio.NewReader(os.Stdin)
570 line, _ := rd.ReadString('\n')
571 line = strings.TrimSpace(line)
572 var n int
573 if _, err := fmt.Sscanf(line, "%d", &n); err != nil || n < 1 || n > len(files) {
574 return "", fmt.Errorf("invalid selection")
575 }
576 return files[n-1], nil
577}
578
579// WriteNewTranscript appends only messages after index `from` from the
580// agent's transcript to the session.
581func WriteNewTranscript(ag *core.Agent, sess *core.Session, from int) {
582 if sess == nil || ag == nil {
583 return
584 }
585 msgs := ag.Messages()
586 for i := from; i < len(msgs); i++ {
587 _ = sess.AppendMessage(msgs[i])
588 }
589 cum := ag.Cost()
590 _ = sess.AppendUsage(cum, cum)
591}
592
593func readAllStdin() (string, error) {
594 fi, err := os.Stdin.Stat()
595 if err != nil {
596 return "", err
597 }
598 if (fi.Mode() & os.ModeCharDevice) != 0 {
599 return "", nil
600 }
601 b, err := io.ReadAll(os.Stdin)
602 return string(b), err
603}
604
605func printModels() {
606 fmt.Println("provider model id context max-out reasoning source name")
607 for _, m := range provider.Active() {
608 reason := " "
609 if m.Reasoning {
610 reason = "✓"
611 }
612 source := m.Source
613 if source == "" {
614 source = "catalog"
615 }
616 if m.Speculative {
617 source = "speculative"
618 }
619 fmt.Printf("%-10s %-30s %8d %8d %s %-11s %s\n",
620 m.Provider, m.ID, m.ContextWindow, m.MaxOutput, reason, source, m.DisplayName)
621 }
622}