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