1package agent
2
3import (
4 "bufio"
5 "context"
6 "encoding/json"
7 "errors"
8 "fmt"
9 "io"
10 "os"
11 "os/signal"
12 "path/filepath"
13 "strings"
14 "sync"
15 "syscall"
16 "time"
17
18 "github.com/patriceckhart/zot/internal/agent/extensions"
19 "github.com/patriceckhart/zot/internal/agent/modes"
20 "github.com/patriceckhart/zot/internal/auth"
21 "github.com/patriceckhart/zot/internal/core"
22 "github.com/patriceckhart/zot/internal/extproto"
23 "github.com/patriceckhart/zot/internal/provider"
24 "github.com/patriceckhart/zot/internal/skills"
25 "github.com/patriceckhart/zot/internal/swarm"
26 "github.com/patriceckhart/zot/internal/tui"
27)
28
29// interactiveExtHooks is a tiny adapter that lets the extension
30// manager call back into the Interactive instance built later in
31// runInteractive. The forward-declared *modes.Interactive is filled
32// in immediately after manager construction.
33type interactiveExtHooks struct {
34 ivPtr **modes.Interactive
35}
36
37func (h *interactiveExtHooks) iv() *modes.Interactive {
38 if h == nil || h.ivPtr == nil {
39 return nil
40 }
41 return *h.ivPtr
42}
43
44func (h *interactiveExtHooks) Notify(extName, level, message string) {
45 if iv := h.iv(); iv != nil {
46 iv.Notify(extName, level, message)
47 }
48}
49func (h *interactiveExtHooks) Submit(text string) {
50 if iv := h.iv(); iv != nil {
51 iv.Submit(text)
52 }
53}
54func (h *interactiveExtHooks) SubmitSlash(text string) {
55 if iv := h.iv(); iv != nil {
56 iv.SubmitSlash(text)
57 }
58}
59func (h *interactiveExtHooks) Insert(text string) {
60 if iv := h.iv(); iv != nil {
61 iv.Insert(text)
62 }
63}
64func (h *interactiveExtHooks) Display(extName, text string) {
65 if iv := h.iv(); iv != nil {
66 iv.Display(extName, text)
67 }
68}
69func (h *interactiveExtHooks) OpenPanel(extName string, spec extproto.PanelSpec) {
70 if iv := h.iv(); iv != nil {
71 iv.OpenPanel(extName, spec)
72 }
73}
74func (h *interactiveExtHooks) UpdatePanel(extName, panelID, title string, lines []string, footer string) {
75 if iv := h.iv(); iv != nil {
76 iv.UpdatePanel(extName, panelID, title, lines, footer)
77 }
78}
79func (h *interactiveExtHooks) ClosePanel(extName, panelID string) {
80 if iv := h.iv(); iv != nil {
81 iv.ClosePanel(extName, panelID)
82 }
83}
84
85// extToolAdapter bridges *extensions.Manager to the
86// ExtensionToolSource interface declared in build.go (kept narrow to
87// avoid a build->extensions import cycle). One adapter instance per
88// run; used at every Resolve point so re-built agents pick up the
89// same set of extension tools.
90type extToolAdapter struct {
91 mgr *extensions.Manager
92}
93
94func (a *extToolAdapter) Tools() []ExtensionToolInfo {
95 infos := a.mgr.Tools()
96 out := make([]ExtensionToolInfo, len(infos))
97 for i, t := range infos {
98 out[i] = ExtensionToolInfo{
99 Extension: t.Extension,
100 Name: t.Name,
101 Description: t.Description,
102 Schema: t.Schema,
103 }
104 }
105 return out
106}
107
108func (a *extToolAdapter) NewExtensionTool(info ExtensionToolInfo) core.Tool {
109 return extensions.NewTool(a.mgr, extensions.ToolInfo{
110 Extension: info.Extension,
111 Name: info.Name,
112 Description: info.Description,
113 Schema: info.Schema,
114 })
115}
116
117// fanoutAgentEvent translates a core.AgentEvent into the wire-format
118// EventFromHost and pushes it through the extension manager. Only
119// the events that have a clear extension-facing meaning are
120// forwarded; internal-only ones (text_delta, tool_progress) are
121// dropped to keep the per-extension stream sane.
122func trimMessagesForResume(msgs []provider.Message, keepTail int) []provider.Message {
123 if keepTail <= 0 || len(msgs) <= keepTail {
124 return provider.RepairOrphanedToolResults(msgs)
125 }
126 var out []provider.Message
127 start := len(msgs) - keepTail
128 // Preserve the synthetic compaction summary when present so an
129 // already-compacted session stays compacted after resume.
130 if len(msgs) > 0 && msgs[0].Meta["compaction"] == "true" && start > 1 {
131 out = append(out, msgs[0])
132 }
133 // Avoid hydrating a tail that starts with orphan tool_result rows;
134 // provider APIs require those to be paired with an earlier tool_use.
135 for start < len(msgs) && msgs[start].Role == provider.RoleTool {
136 start++
137 }
138 out = append(out, msgs[start:]...)
139 return provider.RepairOrphanedToolResults(out)
140}
141
142func fanoutAgentEvent(mgr *extensions.Manager, ev core.AgentEvent) {
143 if mgr == nil {
144 return
145 }
146 switch e := ev.(type) {
147 case core.EvTurnStart:
148 mgr.EmitEvent(extproto.EventFromHost{Event: "turn_start", Step: e.Step})
149 case core.EvToolCall:
150 mgr.EmitEvent(extproto.EventFromHost{
151 Event: "tool_call", ToolID: e.ID, ToolName: e.Name, ToolArgs: e.Args,
152 })
153 case core.EvAssistantMessage:
154 // Concat the visible text portions of the message; binary
155 // blocks (tool_use, etc.) are skipped because subscribers
156 // usually want a string they can grep / display.
157 var text string
158 for _, c := range e.Message.Content {
159 if tb, ok := c.(provider.TextBlock); ok {
160 text += tb.Text
161 }
162 }
163 mgr.EmitEvent(extproto.EventFromHost{Event: "assistant_message", Text: text})
164 case core.EvTurnEnd:
165 ev := extproto.EventFromHost{Event: "turn_end", Stop: string(e.Stop)}
166 if e.Err != nil {
167 ev.Error = e.Err.Error()
168 }
169 mgr.EmitEvent(ev)
170 }
171}
172
173// Run is the top-level entrypoint for the zot binary.
174func Run(rawArgs []string, version string) error {
175 // Subcommand router: `zot bot ...` is handled separately so the
176 // generic flag parser doesn't reject "bot" as a positional arg.
177 if handled, err := runBotCommand(rawArgs, version); handled {
178 return err
179 }
180 if handled, err := runExtCommand(rawArgs); handled {
181 return err
182 }
183 if handled, err := runUpdateCommand(rawArgs, version); handled {
184 return err
185 }
186 // `zot rpc` is shorthand for `zot --rpc` so third-party apps can
187 // spawn the binary with a clean argv. Strip the leading 'rpc'
188 // token and let the rest flow through the normal arg parser.
189 if len(rawArgs) > 0 && rawArgs[0] == "rpc" {
190 rawArgs = append([]string{"--rpc"}, rawArgs[1:]...)
191 }
192
193 args, err := ParseArgs(rawArgs)
194 if err != nil {
195 PrintHelp(version)
196 return err
197 }
198 if args.Help {
199 PrintHelp(version)
200 return nil
201 }
202 if args.Version {
203 fmt.Println("zot", version)
204 return nil
205 }
206 // Model catalog: load any cached discovery data before we inspect
207 // the model list (list-models, print/json, interactive).
208 LoadCachedModels()
209 LoadUserModels()
210
211 // Repair config.json so a stale (provider, model) pair from an
212 // interrupted /model switch can't strand the user with an
213 // "unknown model" error on the first turn. Runs before any UI
214 // renders so the status bar shows the post-repair pair, not the
215 // broken one. Silent on success.
216 ValidateAndRepairConfig()
217
218 if args.ListModels {
219 printModels()
220 return nil
221 }
222
223 ctx := context.Background()
224
225 // Kick an async refresh of the live model catalog. The first run of
226 // zot hits the network; subsequent runs within CacheTTL do nothing.
227 RefreshModelsAsync()
228
229 switch args.Mode {
230 case ModePrint:
231 return runPrintMode(ctx, args, version)
232 case ModeJSON:
233 return runJSONMode(ctx, args, version)
234 case ModeRPC:
235 return runRPCMode(ctx, args, version)
236 case ModeSwarmAgent:
237 return runSwarmAgentMode(ctx, args, version)
238 default:
239 return runInteractive(ctx, args, version)
240 }
241}
242
243// ---- print / json modes: require credentials, run single-shot ----
244
245// nonInteractiveExtHooks is the HostHooks impl used by print / json
246// modes. They have no TUI, so notify / display go to stderr and
247// submit / insert are no-ops (the extension can't steer a
248// single-shot run once it's in flight anyway).
249type nonInteractiveExtHooks struct{}
250
251func (nonInteractiveExtHooks) Notify(ext, level, message string) {
252 fmt.Fprintf(os.Stderr, "[%s] %s: %s\n", ext, level, message)
253}
254func (nonInteractiveExtHooks) Submit(string) {}
255func (nonInteractiveExtHooks) SubmitSlash(string) {}
256func (nonInteractiveExtHooks) Insert(string) {}
257func (nonInteractiveExtHooks) Display(string, string) {}
258func (nonInteractiveExtHooks) OpenPanel(string, extproto.PanelSpec) {}
259func (nonInteractiveExtHooks) UpdatePanel(string, string, string, []string, string) {}
260func (nonInteractiveExtHooks) ClosePanel(string, string) {}
261
262// setupNonInteractiveExtensions loads --ext paths and (unless
263// --no-ext) runs discovery. Returns the manager so the caller can
264// wire tools into the resolved registry, and a cleanup closure to
265// defer. Mirrors the interactive-mode setup minus the TUI hooks.
266func setupNonInteractiveExtensions(ctx context.Context, args Args, r *Resolved, version string) (*extensions.Manager, func()) {
267 extMgr := extensions.New(ZotHome(), r.CWD, version, r.Provider, r.Model, nonInteractiveExtHooks{})
268 for _, e := range extMgr.LoadExplicit(ctx, args.Exts) {
269 fmt.Fprintln(os.Stderr, "extension load:", e)
270 }
271 if !args.NoExt {
272 for _, e := range extMgr.Discover(ctx) {
273 fmt.Fprintln(os.Stderr, "extension load:", e)
274 }
275 }
276 extMgr.WaitForReady(3 * time.Second)
277 r.MergeExtensionTools(&extToolAdapter{mgr: extMgr})
278 extMgr.EmitEvent(extproto.EventFromHost{Event: "session_start"})
279 return extMgr, func() { extMgr.Stop(2 * time.Second) }
280}
281
282// wireNonInteractiveAgentExtHooks installs the same BeforeToolExecute
283// / BeforeTurn / BeforeAssistantMessage / OnEvent hooks the
284// interactive path wires up, so extensions get their normal
285// event-intercept surface in print / json / rpc flows too.
286func wireNonInteractiveAgentExtHooks(ctx context.Context, ag *core.Agent, extMgr *extensions.Manager) {
287 if ag == nil || extMgr == nil {
288 return
289 }
290 ag.BeforeToolExecute = func(call provider.ToolCallBlock) (bool, string, json.RawMessage) {
291 res := extMgr.InterceptToolCall(ctx, call.ID, call.Name, call.Arguments)
292 if res.Block {
293 return false, res.Reason, nil
294 }
295 return true, "", res.ModifiedArgs
296 }
297 ag.BeforeTurn = func(step int) (bool, string) {
298 res := extMgr.InterceptTurnStart(ctx, step)
299 return !res.Block, res.Reason
300 }
301 ag.BeforeAssistantMessage = func(text string) (bool, string, string) {
302 res := extMgr.InterceptAssistantMessage(ctx, text)
303 if res.Block {
304 return false, res.Reason, ""
305 }
306 return true, "", res.ReplaceText
307 }
308 ag.OnEvent = func(ev core.AgentEvent) { fanoutAgentEvent(extMgr, ev) }
309}
310
311func runPrintMode(ctx context.Context, args Args, version string) error {
312 if args.NoYolo {
313 fmt.Fprintln(os.Stderr, "warning: --no-yolo has no effect in print mode (no interactive prompt available); tools will run without confirmation")
314 }
315 r, err := Resolve(args, true)
316 if err != nil {
317 return err
318 }
319 extMgr, stopExt := setupNonInteractiveExtensions(ctx, args, &r, version)
320 defer stopExt()
321
322 ag := r.NewAgent()
323 wireNonInteractiveAgentExtHooks(ctx, ag, extMgr)
324 sess, _ := openOrCreateSession(args, r, ag, version)
325 defer sess.Close()
326
327 prompt := args.Prompt
328 if prompt == "" {
329 piped, _ := readAllStdin()
330 prompt = strings.TrimSpace(piped)
331 }
332 if prompt == "" {
333 return fmt.Errorf("print mode requires a prompt (arg or stdin)")
334 }
335
336 start := len(ag.Messages())
337 err = modes.RunPrint(ctx, ag, prompt, nil, os.Stdout)
338 WriteNewTranscript(ag, sess, start)
339 return err
340}
341
342func runJSONMode(ctx context.Context, args Args, version string) error {
343 if args.NoYolo {
344 fmt.Fprintln(os.Stderr, "warning: --no-yolo has no effect in json mode (no interactive prompt available); tools will run without confirmation")
345 }
346 r, err := Resolve(args, true)
347 if err != nil {
348 return err
349 }
350 extMgr, stopExt := setupNonInteractiveExtensions(ctx, args, &r, version)
351 defer stopExt()
352
353 ag := r.NewAgent()
354 wireNonInteractiveAgentExtHooks(ctx, ag, extMgr)
355 sess, _ := openOrCreateSession(args, r, ag, version)
356 defer sess.Close()
357
358 prompt := args.Prompt
359 if prompt == "" {
360 piped, _ := readAllStdin()
361 prompt = strings.TrimSpace(piped)
362 }
363 if prompt == "" {
364 return fmt.Errorf("json mode requires a prompt (arg or stdin)")
365 }
366
367 start := len(ag.Messages())
368 err = modes.RunJSON(ctx, ag, prompt, nil, os.Stdout)
369 WriteNewTranscript(ag, sess, start)
370 return err
371}
372
373// ---- interactive mode: opens the TUI even without credentials ----
374
375func runInteractive(ctx context.Context, args Args, version string) error {
376 // Resolve WITHOUT requiring credentials.
377 r, err := Resolve(args, false)
378 if err != nil {
379 return err
380 }
381
382 authStore := AuthStoreFor()
383 mgr := auth.NewManager(authStore)
384 defer mgr.Close()
385
386 // Keep the sandbox pointer stable across agent rebuilds (login / model
387 // switch). The Interactive UI toggles the lock via this pointer, and
388 // rebuilt tool instances must share the same one so the lock sticks.
389 sharedSandbox := r.Sandbox
390
391 // Build the extension manager BEFORE the agent so we can fold
392 // extension-defined tools into the registry. Forward-declare iv so
393 // the host hooks adapter can dereference it after construction.
394 var iv *modes.Interactive
395 extHooks := &interactiveExtHooks{ivPtr: &iv}
396 extMgr := extensions.New(ZotHome(), r.CWD, version, r.Provider, r.Model, extHooks)
397 // --ext paths first so they win against installed extensions of
398 // the same name (loadOne's first-write-wins semantics).
399 for _, e := range extMgr.LoadExplicit(ctx, args.Exts) {
400 fmt.Fprintln(os.Stderr, "extension load:", e)
401 }
402 // --no-ext skips the global + project-local discovery scan;
403 // explicit --ext paths above are still honoured so you can run
404 // "only this extension" with --no-ext --ext ./x.
405 if !args.NoExt {
406 for _, e := range extMgr.Discover(ctx) {
407 fmt.Fprintln(os.Stderr, "extension load:", e)
408 }
409 }
410 // Wait briefly for extensions to flush their initial register_tool
411 // frames before we build the agent's tool registry. Half a second
412 // is plenty for any extension that's actually well-behaved; ones
413 // that don't send a ready frame eat the full grace and proceed.
414 // 3s is the per-extension grace period for the ready frame.
415 // Native binaries are instant; runtimes like `npx tsx` take ~1.5s
416 // from cold cache. The wait is tight only for extensions that
417 // haven't sent ready by then; ones that signalled earlier release
418 // the wait immediately.
419 extMgr.WaitForReady(3 * time.Second)
420 defer extMgr.Stop(2 * time.Second)
421
422 extToolAdapter := &extToolAdapter{mgr: extMgr}
423 r.MergeExtensionTools(extToolAdapter)
424
425 // Confirmation gate: when --no-yolo is on, the agent must ask
426 // the user before every tool call. In interactive mode the TUI
427 // provides the Confirmer; in print/json/rpc modes there's no
428 // way to prompt, so the gate is constructed with a nil inner
429 // which auto-refuses every call with a helpful reason.
430 var confirmGate *core.ConfirmGate
431 if args.NoYolo {
432 confirmGate = core.NewConfirmGate(nil) // set below for interactive
433 }
434
435 // Capture current args in a closure so BuildAgent can re-resolve
436 // after a successful login (picks up the newly stored credential).
437 wireAgentExt := func(a *core.Agent) *core.Agent {
438 if a == nil {
439 return a
440 }
441 a.BeforeToolExecute = func(call provider.ToolCallBlock) (bool, string, json.RawMessage) {
442 // Confirm gate runs FIRST: if the user refused, we don't
443 // waste extension-intercept time or let guards see the call.
444 if confirmGate != nil {
445 ok, reason, _ := confirmGate.Check(call.Name, core.BuildPreview(call.Arguments, 120))
446 if !ok {
447 return false, reason, nil
448 }
449 }
450 r := extMgr.InterceptToolCall(ctx, call.ID, call.Name, call.Arguments)
451 if r.Block {
452 return false, r.Reason, nil
453 }
454 return true, "", r.ModifiedArgs
455 }
456 a.BeforeTurn = func(step int) (bool, string) {
457 r := extMgr.InterceptTurnStart(ctx, step)
458 return !r.Block, r.Reason
459 }
460 a.BeforeAssistantMessage = func(text string) (bool, string, string) {
461 r := extMgr.InterceptAssistantMessage(ctx, text)
462 if r.Block {
463 return false, r.Reason, ""
464 }
465 return true, "", r.ReplaceText
466 }
467 a.OnEvent = func(ev core.AgentEvent) { fanoutAgentEvent(extMgr, ev) }
468 return a
469 }
470
471 buildAgent := func() (*core.Agent, string, string, error) {
472 resolved, err := Resolve(args, true)
473 if err != nil {
474 return nil, "", "", err
475 }
476 resolved.UseSandbox(sharedSandbox)
477 resolved.MergeExtensionTools(extToolAdapter)
478 return wireAgentExt(resolved.NewAgent()), resolved.Provider, resolved.Model, nil
479 }
480
481 // Rebuild agent with an explicit provider/model override.
482 buildAgentFor := func(providerOverride, modelOverride string) (*core.Agent, string, string, error) {
483 next := args
484 if providerOverride != "" {
485 next.Provider = providerOverride
486 }
487 if modelOverride != "" {
488 next.Model = modelOverride
489 }
490 resolved, err := Resolve(next, true)
491 if err != nil {
492 return nil, "", "", err
493 }
494 resolved.UseSandbox(sharedSandbox)
495 resolved.MergeExtensionTools(extToolAdapter)
496 return wireAgentExt(resolved.NewAgent()), resolved.Provider, resolved.Model, nil
497 }
498
499 // Rebuild agent for the rescue picker after a recoverable failure.
500 // Unlike buildAgentFor, this drops launch-time --api-key and
501 // --base-url overrides because those are typically the cause of the
502 // rescue (a bad key, a typo'd base URL, or a corporate gateway that
503 // only the originally-picked provider needed). Re-resolving without
504 // them lets the rescue retry use env vars / auth.json / provider
505 // defaults the way zot would have without the overrides.
506 buildAgentForRescue := func(providerOverride, modelOverride string) (*core.Agent, string, string, error) {
507 next := args
508 next.APIKey = ""
509 next.BaseURL = ""
510 if providerOverride != "" {
511 next.Provider = providerOverride
512 }
513 if modelOverride != "" {
514 next.Model = modelOverride
515 }
516 resolved, err := Resolve(next, true)
517 if err != nil {
518 return nil, "", "", err
519 }
520 resolved.UseSandbox(sharedSandbox)
521 resolved.MergeExtensionTools(extToolAdapter)
522 return wireAgentExt(resolved.NewAgent()), resolved.Provider, resolved.Model, nil
523 }
524
525 var ag *core.Agent
526 if r.HasCredential() {
527 ag = wireAgentExt(r.NewAgent())
528 }
529
530 // /reload-ext callback: after the manager has respawned every
531 // extension, re-resolve the tool registry (built-ins + freshly-
532 // registered extension tools) and swap it onto the current
533 // agent in-place. The current agent may have been replaced by a
534 // /model swap since spawn, so re-read the live `ag` on each
535 // invocation.
536 extMgr.SetOnReload(func() {
537 current := ag
538 if current == nil {
539 return
540 }
541 resolved, err := Resolve(args, true)
542 if err != nil {
543 return
544 }
545 resolved.UseSandbox(sharedSandbox)
546 resolved.MergeExtensionTools(extToolAdapter)
547 current.SetTools(resolved.ToolRegistry)
548 })
549
550 // Fire session_start once we know the manager's running.
551 extMgr.EmitEvent(extproto.EventFromHost{Event: "session_start"})
552
553 var sess *core.Session
554 var sessBaselineMsgs int // messages already on disk when current session opened
555 // swarmMgr is constructed below, but loadSession (defined before
556 // the construction site) needs to re-scope it whenever the user
557 // swaps sessions. Forward-declare here so the closure can
558 // reference it; the assignment happens at the construction line.
559 var swarmMgr *swarm.Swarm
560 // persistMu guards sess + sessBaselineMsgs against concurrent access
561 // from the agent loop's per-message persistence hook (runs on the
562 // agent goroutine) and the TUI's session swap / flush callbacks
563 // (run on the TUI goroutine). Without this, a /sessions swap that
564 // races with a finishing turn could double-write or lose messages.
565 var persistMu sync.Mutex
566 if !args.NoSess && ag != nil {
567 sess, _ = openOrCreateSession(args, r, ag, version)
568 if ag != nil {
569 sessBaselineMsgs = len(ag.Messages())
570 }
571 }
572 defer func() {
573 persistMu.Lock()
574 defer persistMu.Unlock()
575 if sess != nil {
576 sess.Close()
577 }
578 }()
579
580 // persistMessage is the per-message hook bound to the agent. It
581 // appends each new transcript message to the live session as soon
582 // as it lands, so a kill / closed terminal / OS crash costs at
583 // most the in-flight turn instead of the whole session. The
584 // baseline counter advances in lock-step so the exit-time flush
585 // doesn't double-write rows already on disk.
586 persistMessage := func(m provider.Message) {
587 persistMu.Lock()
588 defer persistMu.Unlock()
589 if sess == nil {
590 return
591 }
592 if err := sess.AppendMessage(m); err == nil {
593 sessBaselineMsgs++
594 }
595 }
596 persistUsage := func(cum provider.Usage) {
597 persistMu.Lock()
598 defer persistMu.Unlock()
599 if sess == nil {
600 return
601 }
602 _ = sess.AppendUsage(cum, cum)
603 }
604 persistCompaction := func(messages []provider.Message) {
605 persistMu.Lock()
606 defer persistMu.Unlock()
607 if sess == nil {
608 return
609 }
610 if err := sess.AppendCompaction(messages); err == nil {
611 sessBaselineMsgs = len(messages)
612 }
613 }
614 wireAgentPersist := func(a *core.Agent) *core.Agent {
615 if a == nil {
616 return a
617 }
618 a.OnMessageAppended = persistMessage
619 a.OnUsage = persistUsage
620 a.OnTranscriptCompacted = persistCompaction
621 return a
622 }
623 wireAgentPersist(ag)
624
625 // Re-wrap the build closures so any agent constructed by the TUI
626 // (login, /model swap to a different provider) also gets the
627 // persistence hooks. Without this, switching provider would
628 // silently revert to the old in-memory-only behaviour.
629 baseBuildAgent := buildAgent
630 buildAgent = func() (*core.Agent, string, string, error) {
631 a, p, m, err := baseBuildAgent()
632 return wireAgentPersist(a), p, m, err
633 }
634 baseBuildAgentFor := buildAgentFor
635 buildAgentFor = func(providerOverride, modelOverride string) (*core.Agent, string, string, error) {
636 a, p, m, err := baseBuildAgentFor(providerOverride, modelOverride)
637 return wireAgentPersist(a), p, m, err
638 }
639 baseBuildAgentForRescue := buildAgentForRescue
640 buildAgentForRescue = func(providerOverride, modelOverride string) (*core.Agent, string, string, error) {
641 a, p, m, err := baseBuildAgentForRescue(providerOverride, modelOverride)
642 return wireAgentPersist(a), p, m, err
643 }
644
645 // loadSession replaces the current session with the one at path and
646 // hands its messages to the agent. Used by the /sessions picker.
647 loadSession := func(path string) error {
648 currentAg := ag // captured
649 if currentAg == nil {
650 return fmt.Errorf("no agent running; log in first")
651 }
652 newSess, msgs, err := core.OpenSession(path)
653 if err != nil {
654 return err
655 }
656 fullMsgCount := len(msgs)
657 msgs = trimMessagesForResume(msgs, 100)
658 persistMu.Lock()
659 // Flush any unsaved messages to the old session before swapping.
660 // Per-message persistence keeps sessBaselineMsgs current, so
661 // this is a defensive no-op in the common case; it still
662 // matters for the rare race where a turn just finished and
663 // hadn't fired its hook yet.
664 if sess != nil {
665 writeNewTranscriptLocked(currentAg, sess, sessBaselineMsgs)
666 _ = sess.Close()
667 }
668 sess = newSess
669 currentAg.SetMessages(msgs)
670 if cum, last, uerr := core.SessionUsageDetail(path); uerr == nil {
671 currentAg.SeedCost(cum)
672 currentAg.SeedLastTurnUsage(last)
673 }
674 // The live agent only receives a compact resume window, but
675 // the session file remains intact. Keep the persistence
676 // baseline at the original on-disk message count so future
677 // turns append after the full session instead of duplicating
678 // the hydrated tail.
679 sessBaselineMsgs = fullMsgCount
680 persistMu.Unlock()
681 // Re-scope the swarm dashboard to the new session so /swarm
682 // only shows agents this session spawned. swarmMgr may be nil
683 // here if we haven't reached the construction site yet (it
684 // shouldn't be, since the interactive loop is what triggers
685 // loadSession, but the nil check is cheap insurance).
686 if swarmMgr != nil && newSess != nil {
687 swarmMgr.SetActiveSession(newSess.ID)
688 }
689 return nil
690 }
691
692 // changeCWD switches the running session to a new working directory.
693 // Wired into InteractiveConfig.ChangeCWD and invoked by the hidden
694 // /cd slash command (which itself is only fired by the workspaces
695 // extension's panel-key Enter handler today; the user can type /cd
696 // directly but it's not in autocomplete / help / the README).
697 //
698 // Steps, in order:
699 // 1. resolve + validate the new path (~ expansion, abs/rel)
700 // 2. close the current session, flushing pending messages
701 // 3. mutate captured args.CWD + r.CWD so future buildAgent
702 // calls bind to the new cwd
703 // 4. re-root the shared sandbox («·«) so /jail follows the
704 // session into the new cwd instead of widening or silently
705 // dropping
706 // 5. rebuild the agent via buildAgent() so tools, AGENTS.md
707 // addendum, system prompt, sessions dir all bind correctly
708 // 6. open a fresh session in the new cwd's bucket
709 // 7. push the new state into the running Interactive
710 // 8. re-scope the swarm dashboard to the freshly-opened session
711 //
712 // The /jail state is preserved verbatim: if the sandbox was locked
713 // to the old cwd, it stays locked, just re-pointed at the new one.
714 changeCWD := func(path string) error {
715 if path == "" {
716 return fmt.Errorf("empty path")
717 }
718 // ~ expansion.
719 if path == "~" || strings.HasPrefix(path, "~/") {
720 home, herr := os.UserHomeDir()
721 if herr != nil || home == "" {
722 return fmt.Errorf("cannot expand ~: %v", herr)
723 }
724 if path == "~" {
725 path = home
726 } else {
727 path = filepath.Join(home, path[2:])
728 }
729 }
730 if !filepath.IsAbs(path) {
731 path = filepath.Join(args.CWD, path)
732 }
733 absPath, err := filepath.Abs(path)
734 if err != nil {
735 return err
736 }
737 info, err := os.Stat(absPath)
738 if err != nil {
739 return err
740 }
741 if !info.IsDir() {
742 return fmt.Errorf("not a directory: %s", absPath)
743 }
744
745 currentAg := ag
746 if currentAg == nil {
747 return fmt.Errorf("no agent running; log in first")
748 }
749
750 // Close the current session before we drop the reference.
751 // Per-message persistence keeps it current already; this is
752 // a defensive flush + final fsync via Close.
753 persistMu.Lock()
754 if sess != nil {
755 writeNewTranscriptLocked(currentAg, sess, sessBaselineMsgs)
756 _ = sess.Close()
757 sess = nil
758 }
759 sessBaselineMsgs = 0
760 persistMu.Unlock()
761
762 // Mutate captured state so subsequent agent rebuilds and
763 // session opens see the new cwd.
764 wasJailed := sharedSandbox != nil && sharedSandbox.Locked()
765 args.CWD = absPath
766 r.CWD = absPath
767 if sharedSandbox != nil {
768 sharedSandbox.Root = absPath
769 if wasJailed {
770 sharedSandbox.Lock()
771 } else {
772 sharedSandbox.Unlock()
773 }
774 }
775
776 // Rebuild the agent so tools / AGENTS.md / system prompt
777 // re-bind to the new cwd. buildAgent() reads from args + r.
778 newAg, newProvider, newModel, berr := buildAgent()
779 if berr != nil {
780 return fmt.Errorf("rebuild agent: %v", berr)
781 }
782 ag = newAg
783
784 // Fresh session in the new cwd's bucket. We bypass
785 // openOrCreateSession's --continue / --resume branches
786 // because /cd's semantics are "start fresh here", matching
787 // what relaunching `zot --cwd <path>` would do today.
788 if !args.NoSess {
789 core.PruneEmptySessions(ZotHome(), absPath)
790 newSess, serr := core.NewSession(ZotHome(), absPath, newProvider, newModel, version)
791 if serr != nil {
792 return fmt.Errorf("open session in %s: %v", absPath, serr)
793 }
794 persistMu.Lock()
795 sess = newSess
796 sessBaselineMsgs = 0
797 persistMu.Unlock()
798 }
799
800 // Push the new state into the running Interactive.
801 if iv != nil {
802 iv.ApplyChangedCWD(newAg, newProvider, newModel, absPath)
803 }
804
805 // Re-scope the swarm dashboard to the new session.
806 if swarmMgr != nil && sess != nil {
807 swarmMgr.SetActiveSession(sess.ID)
808 }
809 return nil
810 }
811
812 term := tui.NewProcTerm()
813
814 // Kick off the async update check so the banner can appear when the
815 // http response eventually arrives (usually <1s on cached DNS). Map
816 // agent.UpdateInfo -> modes.UpdateInfo here to avoid a cyclic import.
817 updateCh := make(chan modes.UpdateInfo, 1)
818 go func() {
819 defer close(updateCh)
820 src := <-CheckForUpdateAsync(ZotHome(), version)
821 updateCh <- modes.UpdateInfo{
822 Current: src.Current,
823 Latest: src.Latest,
824 Available: src.Available,
825 URL: src.URL,
826 }
827 }()
828
829 // Changelog: when the running version differs from the last
830 // version whose release notes the user dismissed, fetch the
831 // release body from GitHub and have the TUI show it once. On
832 // first-ever launch (no prior LastChangelogShown), seed the
833 // stored version silently — don't dump release notes at someone
834 // who just installed.
835 changelogCh := make(chan modes.ChangelogPayload, 1)
836 go func() {
837 defer close(changelogCh)
838 cfg, _ := LoadConfig()
839 if cfg.LastChangelogShown == "" {
840 SeedChangelogVersion(version)
841 return
842 }
843 if !ShouldShowChangelog(version, cfg) {
844 return
845 }
846 info := <-FetchChangelogAsync(version)
847 if info.Body == "" {
848 return
849 }
850 // For dev builds (0.0.0), skip if the latest release was
851 // already shown (stored by the dismiss callback).
852 if version == "0.0.0" && info.Version == cfg.LastChangelogShown {
853 return
854 }
855 changelogCh <- modes.ChangelogPayload{
856 Version: info.Version,
857 Body: info.Body,
858 URL: info.URL,
859 }
860 }()
861
862 initialCfg, _ := LoadConfig()
863
864 // Build the swarm supervisor. Root lives under ZotHome/swarm so
865 // worktrees survive across zot sessions and the user can hunt
866 // them down with `git worktree list` if anything misbehaves.
867 swarmMgr = swarm.New(swarm.Config{
868 Root: filepath.Join(ZotHome(), "swarm"),
869 RepoRoot: r.CWD,
870 })
871 // Pull any previously-spawned agents off disk so the user can see,
872 // resume, or remove them from the dashboard. Failures here aren't
873 // fatal — the supervisor still works for new agents.
874 _, _ = swarmMgr.Reload()
875 // Scope the dashboard to the active host session so /swarm only
876 // shows agents this session spawned (and any pre-upgrade unscoped
877 // agents — see SnapshotAll docs). Updated again whenever the
878 // user swaps sessions via loadSession below.
879 if sess != nil {
880 swarmMgr.SetActiveSession(sess.ID)
881 }
882 // Best-effort shutdown on interactive exit: stop all running
883 // agents so they don't outlive their parent zot.
884 defer swarmMgr.StopAll()
885
886 iv = modes.NewInteractive(modes.InteractiveConfig{
887 Terminal: term,
888 Theme: tui.DetectThemeFromBackground(80 * time.Millisecond),
889 InlineImagesEnabled: initialCfg.InlineImagesEnabled,
890 SettingsStore: configSettingsStore{},
891 Model: r.Model,
892 Provider: r.Provider,
893 AuthMethod: r.AuthMethod,
894 BaseURL: r.BaseURL,
895 Reasoning: r.Reasoning,
896 SystemPrompt: r.SystemPrompt,
897 Tools: r.ToolRegistry,
898 MaxSteps: r.MaxSteps,
899 CWD: r.CWD,
900 ZotHome: ZotHome(),
901 Version: version,
902 UpdateInfoChan: updateCh,
903 Sandbox: sharedSandbox,
904 Agent: ag,
905 InitialInput: args.Prompt,
906 AuthManager: mgr,
907 BuildAgent: buildAgent,
908 SetKimiCLIFallbackDisabled: SetKimiCLIFallbackDisabled,
909 BuildAgentFor: buildAgentFor,
910 BuildAgentForRescue: buildAgentForRescue,
911 LoggedInProviders: func() []string {
912 var out []string
913 seen := map[string]bool{}
914 for _, p := range knownProviders {
915 if _, _, err := ResolveCredential(p, ""); err == nil && !seen[p] {
916 out = append(out, p)
917 seen[p] = true
918 }
919 }
920 // Ollama models are always available (no auth needed).
921 if !seen["ollama"] {
922 out = append(out, "ollama")
923 }
924 return out
925 },
926 LoadSession: loadSession,
927 ChangeCWD: changeCWD,
928 CurrentSessionPath: func() string {
929 if sess == nil {
930 return ""
931 }
932 return sess.Path
933 },
934 FlushSession: func() {
935 // Append any not-yet-persisted agent messages to the
936 // current session file, then advance the baseline so
937 // the final WriteNewTranscript at exit doesn't write
938 // duplicates. Per-message persistence keeps the on-
939 // disk file current already, so this is mostly a
940 // defensive flush — still needed for /session export
941 // to guarantee the exported bytes include the very
942 // last in-flight turn.
943 currentAg := iv.Agent()
944 if currentAg == nil {
945 return
946 }
947 persistMu.Lock()
948 defer persistMu.Unlock()
949 if sess == nil {
950 return
951 }
952 writeNewTranscriptLocked(currentAg, sess, sessBaselineMsgs)
953 sessBaselineMsgs = len(currentAg.Messages())
954 },
955 Extensions: extMgr,
956 Swarm: swarmMgr,
957 ChangelogChan: changelogCh,
958 OnChangelogDismiss: func() {
959 // For dev builds (0.0.0) store the actual release version
960 // so the same changelog doesn't show again next launch.
961 // For real builds, store the binary version.
962 v := version
963 if v == "0.0.0" {
964 if iv != nil && iv.ChangelogVersion() != "" {
965 v = iv.ChangelogVersion()
966 }
967 }
968 _ = MarkChangelogShown(v)
969 },
970 SkillSnapshot: func() []*skills.Skill {
971 if args.NoSkill {
972 // --no-skill: nothing for the picker to show.
973 return nil
974 }
975 // Re-discover so the picker reflects edits made during
976 // the session. Cheap; SKILL.md files are small. Filter
977 // out built-in skills — they're hidden from user-facing
978 // surfaces because they're implementation detail; the
979 // model still sees them through the system-prompt
980 // manifest and the skill tool. User skills only appear
981 // when --with-skills is set; without it the picker shows
982 // nothing but the model still has the built-ins.
983 userHome, _ := os.UserHomeDir()
984 list, _ := skills.Discover(ZotHome(), r.CWD, userHome, args.WithSkills)
985 return skills.VisibleSkills(list)
986 },
987 NoYolo: args.NoYolo,
988 ConfirmGate: confirmGate,
989 PersistModel: func(providerName, model string) {
990 // Update config.json so next launch uses the same pick.
991 cfg, _ := LoadConfig()
992 cfg.Provider = providerName
993 cfg.Model = model
994 _ = SaveConfig(cfg)
995 // Update the active session's meta so resume picks this up.
996 if sess != nil {
997 _ = sess.UpdateModel(providerName, model)
998 }
999 },
1000 })
1001
1002 // Bind the interactive TUI as the Confirmer. We deferred this
1003 // until now because the gate is constructed before the TUI
1004 // (the BeforeToolExecute closure captures it). SetConfirmer
1005 // is mutex-guarded on the gate so this is safe.
1006 if confirmGate != nil {
1007 confirmGate.SetConfirmer(iv)
1008 }
1009
1010 // Signal-driven flush: a SIGTERM / SIGHUP to the zot process
1011 // (closed terminal window, system shutdown, kill) used to lose
1012 // the entire in-memory transcript because the deferred post-Run
1013 // flush below never ran. Per-message persistence above covers
1014 // most of it; this handler writes any in-flight remainder and
1015 // then exits the process so we don't double-paint over a
1016 // broken terminal that the TUI's restore deferreds can no
1017 // longer fix from a signal context.
1018 //
1019 // SIGINT is intentionally NOT handled here — the TUI consumes
1020 // Ctrl+C as a regular key event for cancel/clear semantics, and
1021 // installing a SIGINT notifier here would swallow it.
1022 sigCh := make(chan os.Signal, 1)
1023 signal.Notify(sigCh, syscall.SIGTERM, syscall.SIGHUP)
1024 defer signal.Stop(sigCh)
1025 go func() {
1026 _, ok := <-sigCh
1027 if !ok {
1028 return
1029 }
1030 if finalAg := iv.Agent(); finalAg != nil {
1031 persistMu.Lock()
1032 if sess != nil {
1033 writeNewTranscriptLocked(finalAg, sess, sessBaselineMsgs)
1034 sessBaselineMsgs = len(finalAg.Messages())
1035 _ = sess.Close()
1036 sess = nil
1037 }
1038 persistMu.Unlock()
1039 }
1040 // Exit cleanly. Re-raising the signal would skip os.Exit's
1041 // at-exit hooks; explicit exit is fine because we've already
1042 // flushed the only at-risk state (the session file).
1043 os.Exit(0)
1044 }()
1045
1046 runErr := iv.Run(ctx)
1047
1048 // Flush final transcript to session (only if we had / ended up with an agent).
1049 if finalAg := iv.Agent(); finalAg != nil {
1050 persistMu.Lock()
1051 if sess != nil {
1052 writeNewTranscriptLocked(finalAg, sess, sessBaselineMsgs)
1053 sessBaselineMsgs = len(finalAg.Messages())
1054 }
1055 persistMu.Unlock()
1056 }
1057 return runErr
1058}
1059
1060// openOrCreateSession returns a session for the run. sess may be nil
1061// with a nil error if session persistence is disabled.
1062func openOrCreateSession(args Args, r Resolved, ag *core.Agent, version string) (*core.Session, error) {
1063 if args.NoSess {
1064 return nil, nil
1065 }
1066 // Sweep meta-only files left over from older zot versions (and from
1067 // any session that crashed before its first AppendMessage). Cheap;
1068 // reads the first few bytes of each file in the cwd's session dir.
1069 core.PruneEmptySessions(ZotHome(), args.CWD)
1070 var (
1071 s *core.Session
1072 msgs []provider.Message
1073 err error
1074 )
1075 switch {
1076 case args.Session != "":
1077 s, msgs, err = core.OpenSession(args.Session)
1078 // The swarm-agent child passes a fixed --session path that
1079 // may not exist yet on first Spawn. Treat ENOENT as "create
1080 // a fresh session AT THIS PATH" so the conversation actually
1081 // gets persisted; without this fallback the swarm child runs
1082 // with sess==nil and every Resume re-starts with no memory
1083 // of the prior turns. Other openers (--continue / --resume /
1084 // the picker) never see ENOENT here because they only choose
1085 // paths that already exist on disk.
1086 if err != nil && errors.Is(err, os.ErrNotExist) {
1087 s, err = core.NewSessionAtPath(args.Session, args.CWD, r.Provider, r.Model, version)
1088 msgs = nil
1089 }
1090 case args.Continue:
1091 latest := core.LatestSession(ZotHome(), args.CWD)
1092 if latest != "" {
1093 s, msgs, err = core.OpenSession(latest)
1094 }
1095 case args.Resume:
1096 picked, perr := pickSession(args.CWD)
1097 if perr != nil {
1098 return nil, perr
1099 }
1100 if picked != "" {
1101 s, msgs, err = core.OpenSession(picked)
1102 }
1103 }
1104 if err != nil {
1105 return nil, err
1106 }
1107 if s != nil {
1108 ag.SetMessages(msgs)
1109 if cum, last, uerr := core.SessionUsageDetail(s.Path); uerr == nil {
1110 ag.SeedCost(cum)
1111 ag.SeedLastTurnUsage(last)
1112 }
1113 return s, nil
1114 }
1115 return core.NewSession(ZotHome(), args.CWD, r.Provider, r.Model, version)
1116}
1117
1118func pickSession(cwd string) (string, error) {
1119 files := core.ListSessions(ZotHome(), cwd)
1120 if len(files) == 0 {
1121 fmt.Fprintln(os.Stderr, "no sessions for", cwd)
1122 return "", nil
1123 }
1124 for i, f := range files {
1125 fmt.Fprintf(os.Stderr, " %2d) %s\n", i+1, f)
1126 }
1127 fmt.Fprint(os.Stderr, "pick #: ")
1128 rd := bufio.NewReader(os.Stdin)
1129 line, _ := rd.ReadString('\n')
1130 line = strings.TrimSpace(line)
1131 var n int
1132 if _, err := fmt.Sscanf(line, "%d", &n); err != nil || n < 1 || n > len(files) {
1133 return "", fmt.Errorf("invalid selection")
1134 }
1135 return files[n-1], nil
1136}
1137
1138// WriteNewTranscript appends only messages after index `from` from the
1139// agent's transcript to the session. Used by callers that don't hold
1140// the persistMu (non-interactive print/json modes which run a single
1141// turn under their own goroutine).
1142func WriteNewTranscript(ag *core.Agent, sess *core.Session, from int) {
1143 writeNewTranscriptLocked(ag, sess, from)
1144}
1145
1146// writeNewTranscriptLocked is the same as WriteNewTranscript. The
1147// suffix marks that interactive callers must hold persistMu when
1148// invoking it so concurrent appends from the agent loop don't race
1149// with this catch-up flush.
1150func writeNewTranscriptLocked(ag *core.Agent, sess *core.Session, from int) {
1151 if sess == nil || ag == nil {
1152 return
1153 }
1154 msgs := ag.Messages()
1155 for i := from; i < len(msgs); i++ {
1156 _ = sess.AppendMessage(msgs[i])
1157 }
1158 cum := ag.Cost()
1159 _ = sess.AppendUsage(cum, cum)
1160}
1161
1162func readAllStdin() (string, error) {
1163 fi, err := os.Stdin.Stat()
1164 if err != nil {
1165 return "", err
1166 }
1167 if (fi.Mode() & os.ModeCharDevice) != 0 {
1168 return "", nil
1169 }
1170 b, err := io.ReadAll(os.Stdin)
1171 return string(b), err
1172}
1173
1174func printModels() {
1175 models := provider.Active()
1176
1177 // Compute column widths from actual data so wide providers (e.g.
1178 // xiaomi-token-plan-sgp) and long bedrock model ids don't force the
1179 // `name` column off-screen. Floors mirror the historical layout so
1180 // short catalogs look the same as before.
1181 provW, idW, srcW := len("provider"), len("model id"), len("source")
1182 for _, m := range models {
1183 if w := len(m.Provider); w > provW {
1184 provW = w
1185 }
1186 if w := len(m.ID); w > idW {
1187 idW = w
1188 }
1189 source := m.Source
1190 if source == "" {
1191 source = "catalog"
1192 }
1193 if m.Speculative {
1194 source = "speculative"
1195 }
1196 if w := len(source); w > srcW {
1197 srcW = w
1198 }
1199 }
1200
1201 header := fmt.Sprintf("%-*s %-*s %8s %8s %s %-*s %s",
1202 provW, "provider",
1203 idW, "model id",
1204 "context", "max-out", "reasoning",
1205 srcW, "source",
1206 "name")
1207 fmt.Println(header)
1208
1209 for _, m := range models {
1210 reason := " "
1211 if m.Reasoning {
1212 reason = "✓"
1213 }
1214 source := m.Source
1215 if source == "" {
1216 source = "catalog"
1217 }
1218 if m.Speculative {
1219 source = "speculative"
1220 }
1221 fmt.Printf("%-*s %-*s %8d %8d %s %-*s %s\n",
1222 provW, m.Provider,
1223 idW, m.ID,
1224 m.ContextWindow, m.MaxOutput,
1225 reason,
1226 srcW, source,
1227 m.DisplayName)
1228 }
1229}