1package agent
2
3import (
4 "bufio"
5 "context"
6 "fmt"
7 "io"
8 "os"
9 "strings"
10
11 "github.com/patriceckhart/zot/internal/agent/modes"
12 "github.com/patriceckhart/zot/internal/auth"
13 "github.com/patriceckhart/zot/internal/core"
14 "github.com/patriceckhart/zot/internal/provider"
15 "github.com/patriceckhart/zot/internal/tui"
16)
17
18// Run is the top-level entrypoint for the zot binary.
19func Run(rawArgs []string, version string) error {
20 // Subcommand router: `zot bot ...` is handled separately so the
21 // generic flag parser doesn't reject "bot" as a positional arg.
22 if handled, err := runBotCommand(rawArgs, version); handled {
23 return err
24 }
25
26 args, err := ParseArgs(rawArgs)
27 if err != nil {
28 PrintHelp(version)
29 return err
30 }
31 if args.Help {
32 PrintHelp(version)
33 return nil
34 }
35 if args.Version {
36 fmt.Println("zot", version)
37 return nil
38 }
39 // Model catalog: load any cached discovery data before we inspect
40 // the model list (list-models, print/json, interactive).
41 LoadCachedModels()
42
43 if args.ListModels {
44 printModels()
45 return nil
46 }
47
48 ctx := context.Background()
49
50 // Kick an async refresh of the live model catalog. The first run of
51 // zot hits the network; subsequent runs within CacheTTL do nothing.
52 RefreshModelsAsync()
53
54 switch args.Mode {
55 case ModePrint:
56 return runPrintMode(ctx, args, version)
57 case ModeJSON:
58 return runJSONMode(ctx, args, version)
59 default:
60 return runInteractive(ctx, args, version)
61 }
62}
63
64// ---- print / json modes: require credentials, run single-shot ----
65
66func runPrintMode(ctx context.Context, args Args, version string) error {
67 r, err := Resolve(args, true)
68 if err != nil {
69 return err
70 }
71 ag := r.NewAgent()
72 sess, _ := openOrCreateSession(args, r, ag, version)
73 defer sess.Close()
74
75 prompt := args.Prompt
76 if prompt == "" {
77 piped, _ := readAllStdin()
78 prompt = strings.TrimSpace(piped)
79 }
80 if prompt == "" {
81 return fmt.Errorf("print mode requires a prompt (arg or stdin)")
82 }
83
84 start := len(ag.Messages())
85 err = modes.RunPrint(ctx, ag, prompt, nil, os.Stdout)
86 WriteNewTranscript(ag, sess, start)
87 return err
88}
89
90func runJSONMode(ctx context.Context, args Args, version string) error {
91 r, err := Resolve(args, true)
92 if err != nil {
93 return err
94 }
95 ag := r.NewAgent()
96 sess, _ := openOrCreateSession(args, r, ag, version)
97 defer sess.Close()
98
99 prompt := args.Prompt
100 if prompt == "" {
101 piped, _ := readAllStdin()
102 prompt = strings.TrimSpace(piped)
103 }
104 if prompt == "" {
105 return fmt.Errorf("json mode requires a prompt (arg or stdin)")
106 }
107
108 start := len(ag.Messages())
109 err = modes.RunJSON(ctx, ag, prompt, nil, os.Stdout)
110 WriteNewTranscript(ag, sess, start)
111 return err
112}
113
114// ---- interactive mode: opens the TUI even without credentials ----
115
116func runInteractive(ctx context.Context, args Args, version string) error {
117 // Resolve WITHOUT requiring credentials.
118 r, err := Resolve(args, false)
119 if err != nil {
120 return err
121 }
122
123 authStore := AuthStoreFor()
124 mgr := auth.NewManager(authStore)
125 defer mgr.Close()
126
127 // Keep the sandbox pointer stable across agent rebuilds (login / model
128 // switch). The Interactive UI toggles the lock via this pointer, and
129 // rebuilt tool instances must share the same one so the lock sticks.
130 sharedSandbox := r.Sandbox
131
132 // Capture current args in a closure so BuildAgent can re-resolve
133 // after a successful login (picks up the newly stored credential).
134 buildAgent := func() (*core.Agent, string, string, error) {
135 resolved, err := Resolve(args, true)
136 if err != nil {
137 return nil, "", "", err
138 }
139 resolved.UseSandbox(sharedSandbox)
140 return resolved.NewAgent(), resolved.Provider, resolved.Model, nil
141 }
142
143 // Rebuild agent with an explicit provider/model override.
144 buildAgentFor := func(providerOverride, modelOverride string) (*core.Agent, string, string, error) {
145 next := args
146 if providerOverride != "" {
147 next.Provider = providerOverride
148 }
149 if modelOverride != "" {
150 next.Model = modelOverride
151 }
152 resolved, err := Resolve(next, true)
153 if err != nil {
154 return nil, "", "", err
155 }
156 resolved.UseSandbox(sharedSandbox)
157 return resolved.NewAgent(), resolved.Provider, resolved.Model, nil
158 }
159
160 var ag *core.Agent
161 if r.HasCredential() {
162 ag = r.NewAgent()
163 }
164
165 var sess *core.Session
166 var sessBaselineMsgs int // messages already on disk when current session opened
167 if !args.NoSess && ag != nil {
168 sess, _ = openOrCreateSession(args, r, ag, version)
169 if ag != nil {
170 sessBaselineMsgs = len(ag.Messages())
171 }
172 }
173 defer func() {
174 if sess != nil {
175 sess.Close()
176 }
177 }()
178
179 // loadSession replaces the current session with the one at path and
180 // hands its messages to the agent. Used by the /sessions picker.
181 loadSession := func(path string) error {
182 currentAg := ag // captured
183 if currentAg == nil {
184 return fmt.Errorf("no agent running; log in first")
185 }
186 newSess, msgs, err := core.OpenSession(path)
187 if err != nil {
188 return err
189 }
190 // Flush any unsaved messages to the old session before swapping.
191 if sess != nil {
192 WriteNewTranscript(currentAg, sess, sessBaselineMsgs)
193 _ = sess.Close()
194 }
195 sess = newSess
196 currentAg.SetMessages(msgs)
197 sessBaselineMsgs = len(msgs)
198 return nil
199 }
200
201 term := tui.NewProcTerm()
202
203 // Kick off the async update check so the banner can appear when the
204 // http response eventually arrives (usually <1s on cached DNS). Map
205 // agent.UpdateInfo -> modes.UpdateInfo here to avoid a cyclic import.
206 updateCh := make(chan modes.UpdateInfo, 1)
207 go func() {
208 defer close(updateCh)
209 src := <-CheckForUpdateAsync(ZotHome(), version)
210 updateCh <- modes.UpdateInfo{
211 Current: src.Current,
212 Latest: src.Latest,
213 Available: src.Available,
214 URL: src.URL,
215 }
216 }()
217
218 iv := modes.NewInteractive(modes.InteractiveConfig{
219 Terminal: term,
220 Theme: tui.Dark,
221 Model: r.Model,
222 Provider: r.Provider,
223 AuthMethod: r.AuthMethod,
224 BaseURL: r.BaseURL,
225 Reasoning: r.Reasoning,
226 SystemPrompt: r.SystemPrompt,
227 Tools: r.ToolRegistry,
228 MaxSteps: r.MaxSteps,
229 CWD: r.CWD,
230 ZotHome: ZotHome(),
231 Version: version,
232 UpdateInfoChan: updateCh,
233 Sandbox: sharedSandbox,
234 Agent: ag,
235 InitialInput: args.Prompt,
236 AuthManager: mgr,
237 BuildAgent: buildAgent,
238 BuildAgentFor: buildAgentFor,
239 LoadSession: loadSession,
240 PersistModel: func(providerName, model string) {
241 // Update config.json so next launch uses the same pick.
242 cfg, _ := LoadConfig()
243 cfg.Provider = providerName
244 cfg.Model = model
245 _ = SaveConfig(cfg)
246 // Update the active session's meta so resume picks this up.
247 if sess != nil {
248 _ = sess.UpdateModel(providerName, model)
249 }
250 },
251 })
252
253 runErr := iv.Run(ctx)
254
255 // Flush final transcript to session (only if we had / ended up with an agent).
256 if finalAg := iv.Agent(); finalAg != nil && sess != nil {
257 WriteNewTranscript(finalAg, sess, sessBaselineMsgs)
258 }
259 return runErr
260}
261
262// openOrCreateSession returns a session for the run. sess may be nil
263// with a nil error if session persistence is disabled.
264func openOrCreateSession(args Args, r Resolved, ag *core.Agent, version string) (*core.Session, error) {
265 if args.NoSess {
266 return nil, nil
267 }
268 // Sweep meta-only files left over from older zot versions (and from
269 // any session that crashed before its first AppendMessage). Cheap;
270 // reads the first few bytes of each file in the cwd's session dir.
271 core.PruneEmptySessions(ZotHome(), args.CWD)
272 var (
273 s *core.Session
274 msgs []provider.Message
275 err error
276 )
277 switch {
278 case args.Session != "":
279 s, msgs, err = core.OpenSession(args.Session)
280 case args.Continue:
281 latest := core.LatestSession(ZotHome(), args.CWD)
282 if latest != "" {
283 s, msgs, err = core.OpenSession(latest)
284 }
285 case args.Resume:
286 picked, perr := pickSession(args.CWD)
287 if perr != nil {
288 return nil, perr
289 }
290 if picked != "" {
291 s, msgs, err = core.OpenSession(picked)
292 }
293 }
294 if err != nil {
295 return nil, err
296 }
297 if s != nil {
298 ag.SetMessages(msgs)
299 return s, nil
300 }
301 return core.NewSession(ZotHome(), args.CWD, r.Provider, r.Model, version)
302}
303
304func pickSession(cwd string) (string, error) {
305 files := core.ListSessions(ZotHome(), cwd)
306 if len(files) == 0 {
307 fmt.Fprintln(os.Stderr, "no sessions for", cwd)
308 return "", nil
309 }
310 for i, f := range files {
311 fmt.Fprintf(os.Stderr, " %2d) %s\n", i+1, f)
312 }
313 fmt.Fprint(os.Stderr, "pick #: ")
314 rd := bufio.NewReader(os.Stdin)
315 line, _ := rd.ReadString('\n')
316 line = strings.TrimSpace(line)
317 var n int
318 if _, err := fmt.Sscanf(line, "%d", &n); err != nil || n < 1 || n > len(files) {
319 return "", fmt.Errorf("invalid selection")
320 }
321 return files[n-1], nil
322}
323
324// WriteNewTranscript appends only messages after index `from` from the
325// agent's transcript to the session.
326func WriteNewTranscript(ag *core.Agent, sess *core.Session, from int) {
327 if sess == nil || ag == nil {
328 return
329 }
330 msgs := ag.Messages()
331 for i := from; i < len(msgs); i++ {
332 _ = sess.AppendMessage(msgs[i])
333 }
334 cum := ag.Cost()
335 _ = sess.AppendUsage(cum, cum)
336}
337
338func readAllStdin() (string, error) {
339 fi, err := os.Stdin.Stat()
340 if err != nil {
341 return "", err
342 }
343 if (fi.Mode() & os.ModeCharDevice) != 0 {
344 return "", nil
345 }
346 b, err := io.ReadAll(os.Stdin)
347 return string(b), err
348}
349
350func printModels() {
351 fmt.Println("provider model id context max-out reasoning source name")
352 for _, m := range provider.Active() {
353 reason := " "
354 if m.Reasoning {
355 reason = "✓"
356 }
357 source := m.Source
358 if source == "" {
359 source = "catalog"
360 }
361 if m.Speculative {
362 source = "speculative"
363 }
364 fmt.Printf("%-10s %-30s %8d %8d %s %-11s %s\n",
365 m.Provider, m.ID, m.ContextWindow, m.MaxOutput, reason, source, m.DisplayName)
366 }
367}