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 iv := modes.NewInteractive(modes.InteractiveConfig{
203 Terminal: term,
204 Theme: tui.Dark,
205 Model: r.Model,
206 Provider: r.Provider,
207 AuthMethod: r.AuthMethod,
208 BaseURL: r.BaseURL,
209 Reasoning: r.Reasoning,
210 SystemPrompt: r.SystemPrompt,
211 Tools: r.ToolRegistry,
212 MaxSteps: r.MaxSteps,
213 CWD: r.CWD,
214 ZotHome: ZotHome(),
215 Sandbox: sharedSandbox,
216 Agent: ag,
217 InitialInput: args.Prompt,
218 AuthManager: mgr,
219 BuildAgent: buildAgent,
220 BuildAgentFor: buildAgentFor,
221 LoadSession: loadSession,
222 PersistModel: func(providerName, model string) {
223 // Update config.json so next launch uses the same pick.
224 cfg, _ := LoadConfig()
225 cfg.Provider = providerName
226 cfg.Model = model
227 _ = SaveConfig(cfg)
228 // Update the active session's meta so resume picks this up.
229 if sess != nil {
230 _ = sess.UpdateModel(providerName, model)
231 }
232 },
233 })
234
235 runErr := iv.Run(ctx)
236
237 // Flush final transcript to session (only if we had / ended up with an agent).
238 if finalAg := iv.Agent(); finalAg != nil && sess != nil {
239 WriteNewTranscript(finalAg, sess, sessBaselineMsgs)
240 }
241 return runErr
242}
243
244// openOrCreateSession returns a session for the run. sess may be nil
245// with a nil error if session persistence is disabled.
246func openOrCreateSession(args Args, r Resolved, ag *core.Agent, version string) (*core.Session, error) {
247 if args.NoSess {
248 return nil, nil
249 }
250 var (
251 s *core.Session
252 msgs []provider.Message
253 err error
254 )
255 switch {
256 case args.Session != "":
257 s, msgs, err = core.OpenSession(args.Session)
258 case args.Continue:
259 latest := core.LatestSession(ZotHome(), args.CWD)
260 if latest != "" {
261 s, msgs, err = core.OpenSession(latest)
262 }
263 case args.Resume:
264 picked, perr := pickSession(args.CWD)
265 if perr != nil {
266 return nil, perr
267 }
268 if picked != "" {
269 s, msgs, err = core.OpenSession(picked)
270 }
271 }
272 if err != nil {
273 return nil, err
274 }
275 if s != nil {
276 ag.SetMessages(msgs)
277 return s, nil
278 }
279 return core.NewSession(ZotHome(), args.CWD, r.Provider, r.Model, version)
280}
281
282func pickSession(cwd string) (string, error) {
283 files := core.ListSessions(ZotHome(), cwd)
284 if len(files) == 0 {
285 fmt.Fprintln(os.Stderr, "no sessions for", cwd)
286 return "", nil
287 }
288 for i, f := range files {
289 fmt.Fprintf(os.Stderr, " %2d) %s\n", i+1, f)
290 }
291 fmt.Fprint(os.Stderr, "pick #: ")
292 rd := bufio.NewReader(os.Stdin)
293 line, _ := rd.ReadString('\n')
294 line = strings.TrimSpace(line)
295 var n int
296 if _, err := fmt.Sscanf(line, "%d", &n); err != nil || n < 1 || n > len(files) {
297 return "", fmt.Errorf("invalid selection")
298 }
299 return files[n-1], nil
300}
301
302// WriteNewTranscript appends only messages after index `from` from the
303// agent's transcript to the session.
304func WriteNewTranscript(ag *core.Agent, sess *core.Session, from int) {
305 if sess == nil || ag == nil {
306 return
307 }
308 msgs := ag.Messages()
309 for i := from; i < len(msgs); i++ {
310 _ = sess.AppendMessage(msgs[i])
311 }
312 cum := ag.Cost()
313 _ = sess.AppendUsage(cum, cum)
314}
315
316func readAllStdin() (string, error) {
317 fi, err := os.Stdin.Stat()
318 if err != nil {
319 return "", err
320 }
321 if (fi.Mode() & os.ModeCharDevice) != 0 {
322 return "", nil
323 }
324 b, err := io.ReadAll(os.Stdin)
325 return string(b), err
326}
327
328func printModels() {
329 fmt.Println("provider model id context max-out reasoning source name")
330 for _, m := range provider.Active() {
331 reason := " "
332 if m.Reasoning {
333 reason = "✓"
334 }
335 source := m.Source
336 if source == "" {
337 source = "catalog"
338 }
339 if m.Speculative {
340 source = "speculative"
341 }
342 fmt.Printf("%-10s %-30s %8d %8d %s %-11s %s\n",
343 m.Provider, m.ID, m.ContextWindow, m.MaxOutput, reason, source, m.DisplayName)
344 }
345}