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