Function bodies 31 total

func (d *Decoder) readLine() (string, error) {
	if d.next != nil {
		// return peeked line if any
		line := *d.next
		d.next = nil
		return line, nil
	}
	for {
		line, err := d.lines.ReadLine()
		// skip comments
		if err != nil || !strings.HasPrefix(line, "#") {
			return line, err
		}
	}
}

func (d *Decoder) Decode(v any, fieldOverrides ...string) error {
	overrides := map[string]string{}
	for _, spec := range fieldOverrides {
		if i := strings.LastIndex(spec, ":"); i >= 0 {
			if spec[i+1:] == "base64" {
				overrides[spec[:i]] = "base64"
			}
		}
	}

	line, err := d.readLine()
	if err != nil {
		return err
	}

	intermediate := map[string]any{}
	if err := json.Unmarshal([]byte(line), &intermediate); err != nil {
		return err
	}

	for {
		line, err := d.readLine()
		if err != nil {
			if errors.Is(err, io.EOF) {
				break
			}
			return err
		}
		if !strings.HasPrefix(line, ".") {
			d.next = &line
			break
		}
		if err := d.readField(intermediate, line, overrides); err != nil {
			return err
		}
	}

	// TODO: like Encoder.Encode, this is gross. For performance we need to
	// fix this.
	encoded, err := json.Marshal(intermediate)
	if err != nil {
		return err
	}
	return json.Unmarshal(encoded, v)
}

func (d *Decoder) readField(m map[string]any, header string, overrides map[string]string) error {
	parts := strings.SplitN(header[1:], "=", 2)
	if len(parts) != 2 {
		return fmt.Errorf("invalid field line: %s", header)
	}
	field, sep := strings.TrimSpace(parts[0]), strings.TrimSpace(parts[1])
	if !strings.HasPrefix(sep, "<<") {
		return fmt.Errorf("invalid field line: %s", header)
	}
	sep = strings.TrimSpace(strings.TrimPrefix(sep, "<<"))

	var buf strings.Builder
	for {
		line, err := d.lines.ReadLine()
		if err != nil {
			return fmt.Errorf("reading value for %s: %w", field, err)
		}
		buf.WriteString(line)
		if s := buf.String(); strings.HasSuffix(s, sep+"\n") {
			val := s[:len(s)-len(sep)-1]
			if overrides[field] == "base64" {
				val = base64.StdEncoding.EncodeToString([]byte(val))
			}
			return d.setField(m, field, val)
		}
	}
}

func (d *Decoder) setField(m map[string]any, field, val string) error {
	parts := strings.Split(field, ".")
	obj := m
	for len(parts) > 1 {
		var sub map[string]any
		lookup, ok := obj[parts[0]]
		if ok {
			sub, ok = lookup.(map[string]any)
			if !ok {
				return fmt.Errorf("invalid redefinition of %q", field)
			}
		} else {
			sub = map[string]any{}
			obj[parts[0]] = sub
		}
		obj = sub
		parts = parts[1:]
	}
	if _, exists := obj[parts[0]]; exists {
		return fmt.Errorf("redefinition of %q", field)
	}
	obj[parts[0]] = val
	return nil
}

func (e *Encoder) Encode(v any, heredocFields ...string) error {

	// TODO: this is gross, to encode and then parse JSON to get normalized
	// fields, but it's easy to reason about. For performance we need to fix
	// this.
	encoded, err := json.Marshal(v)
	if err != nil {
		return err
	}
	intermediate := map[string]any{}
	if err = json.Unmarshal(encoded, &intermediate); err != nil {
		return err
	}

	heredocs := map[string]string{}

	for _, spec := range heredocFields {
		field, encoding := spec, ""
		if i := strings.LastIndex(spec, ":"); i >= 0 {
			if spec[i+1:] == "base64" {
				encoding = "base64"
				field = spec[:i]
			}
		}

		parts := strings.Split(field, ".")
		obj := intermediate

		// traverse parent objects to get to the direct owning object
		for len(parts) > 1 {
			if subobj, ok := obj[parts[0]].(map[string]any); ok {
				obj = subobj
				parts = parts[1:]
			} else {
				break
			}
		}

		if len(parts) == 1 {
			if val, ok := obj[parts[0]].(string); ok {
				content := v

func (e *Encoder) write(obj map[string]any, heredocs map[string]string) error {
	if err := json.NewEncoder(e.w).Encode(obj); err != nil {
		return err
	}
	for field, val := range heredocs {
		sep := e.newSep(val)
		_, err := fmt.Fprintf(e.w, ".%s = <<%s\n%s%s\n", field, sep, val, sep)
		if err != nil {
			return err
		}
	}
	return nil
}

func (e *Encoder) newSep(val string) string {
	i := 0
	for {
		sep := fmt.Sprintf("END%d", i)
		if !strings.Contains(val, sep) {
			return sep
		}
		i++
	}
}

func usage() {
	_, _ = fmt.Fprintf(os.Stderr, "Usage: %s [flags] [session.hjl]\n", os.Args[0])
	_, _ = fmt.Fprintf(os.Stderr, "If no session file is provided, one is created in ~/.ajent/sessions/\n")
	flag.PrintDefaults()
	os.Exit(1)
}

func buildTools(braveAPIKey, searchURL string) []tools.Tool {
	t := []tools.Tool{
		atools.WebFetchTool,
		atools.ReadFileTool,
		atools.ListDirTool,
		atools.EditFileTool,
		atools.BashTool,
		atools.CreateFileTool,
		atools.GrepFileTool,
		atools.TreeTool,
		atools.FindReplaceTool,
	}
	if braveAPIKey != "" {
		t = append(t, atools.NewWebSearchTool(braveAPIKey, searchURL))
	}
	return t
}

func main() {
	flag.Parse()
	ctx := context.Background()

	sessionPath := flag.Arg(0)
	if sessionPath == "" {
		var err error
		sessionPath, err = defaultSessionPath()
		if err != nil {
			fmt.Fprintf(os.Stderr, "Error creating session path: %v\n", err)
			os.Exit(1)
		}
		fmt.Fprintf(os.Stderr, "Session: %s\n", sessionPath)
	}

	var client gen.Gen
	switch strings.ToLower(*flagProvider) {
	case "anthropic":
		client = anthropic.New(*flagAPIKey)
	case "ollama":
		client = ollama.New(*flagURL)
	case "openai":
		c := openai.New(*flagAPIKey)
		if *flagURL != "" {
			c.SetBaseURL(*flagURL)
		}
		client = c
	case "vertexai":
		var err error
		client, err = vertexai.New(vertexai.GoogleConfig{
			Project:    *flagProject,
			Region:     *flagRegion,
			Credential: *flagCredential,
		})
		if err != nil {
			panic(err)
		}
	default:
		usage()
	}

	cfg := Config{
		MaxTokens:  *flagMaxTokens,
		Tools:      buildTools(*flagBraveAPIKey, *flagSearchURL),
		Serializer: NewFileSerializer(sessionPath),

func NewUnbufferedLineReader(r io.Reader, maxLength int) *UnbufferedLineReader {
	return &UnbufferedLineReader{
		r:         r,
		maxLength: maxLength,
	}
}

func (r *UnbufferedLineReader) ReadLine() (rv string, err error) {
	if r.r == nil {
		return "", io.EOF
	}

	var outbuf [4096]byte
	out := outbuf[:0]
	for {
		var b [1]byte
		_, err := io.ReadFull(r.r, b[:])
		if err != nil {
			if errors.Is(err, io.EOF) {
				r.r = nil
				if len(out) > 0 {
					return string(out), nil
				}
			}
			return "", err
		}
		out = append(out, b[0])
		if r.maxLength > 0 && len(out) > r.maxLength {
			return "", errors.New("max line length exceeded")
		}
		if b[0] == '\n' {
			return string(out), nil
		}
	}
}

func (s *FileSerializer) CreateOrOpen(meta SessionMeta) (SerializedSession, SessionMeta, []prompt.Prompt, error) {
	fh, err := os.Open(s.path)
	if err != nil {
		if os.IsNotExist(err) {
			return s.create(meta)
		}
		return nil, SessionMeta{}, nil, err
	}
	defer fh.Close()

	d := hjl.NewDecoder(fh)

	var fileMeta SessionMeta
	if err := d.Decode(&fileMeta); err != nil {
		return nil, SessionMeta{}, nil, err
	}
	var history []prompt.Prompt
	for {
		var p prompt.Prompt
		if err := d.Decode(&p, "tool_call.arguments:base64"); err != nil {
			if errors.Is(err, io.EOF) {
				break
			}
			return nil, SessionMeta{}, nil, err
		}
		history = append(history, p)
	}

	// Reopen the file in append mode for future writes.
	afh, err := os.OpenFile(s.path, os.O_WRONLY|os.O_APPEND, 0644)
	if err != nil {
		return nil, SessionMeta{}, nil, err
	}

	return &fileSession{fh: afh, enc: hjl.NewEncoder(afh)}, fileMeta, history, nil
}

func (s *FileSerializer) create(meta SessionMeta) (SerializedSession, SessionMeta, []prompt.Prompt, error) {
	fh, err := os.Create(s.path)
	if err != nil {
		return nil, SessionMeta{}, nil, err
	}
	enc := hjl.NewEncoder(fh)
	if err := enc.Encode(meta, "system_prompt"); err != nil {
		fh.Close()
		return nil, SessionMeta{}, nil, err
	}
	return &fileSession{fh: fh, enc: enc}, meta, nil, nil
}

func (s *fileSession) Append(prompts ...prompt.Prompt) error {
	for _, p := range prompts {
		if err := s.enc.Encode(p, "text", "tool_response.content", "tool_call.arguments:base64"); err != nil {
			return err
		}
	}
	return s.fh.Sync()
}

func (s *fileSession) Close() error {
	if s.fh != nil {
		return s.fh.Close()
	}
	return nil
}

func NewSession(client gen.Gen, model string,
	input io.Reader, output io.Writer, cfg Config) (*Session, error) {

	opts := []gen.Option{
		gen.WithModel(gen.Model{Provider: client.Provider(), Name: model}),
	}
	if cfg.MaxTokens != 0 {
		opts = append(opts, gen.WithMaxTokens(cfg.MaxTokens))
	}
	if len(cfg.Tools) > 0 {
		opts = append(opts, gen.WithTools(cfg.Tools...))
	}

	var history []prompt.Prompt
	var serialized SerializedSession
	if cfg.Serializer != nil {
		s, meta, loaded, err := cfg.Serializer.CreateOrOpen(
			SessionMeta{SystemPrompt: cfg.SystemPrompt})
		if err != nil {
			return nil, err
		}
		serialized = s
		cfg.SystemPrompt = meta.SystemPrompt
		history = loaded
	}

	if cfg.SystemPrompt != "" {
		opts = append(opts, gen.WithSystem(cfg.SystemPrompt))
	}

	return &Session{
		gen:        client.Generator(opts...),
		input:      private.NewUnbufferedLineReader(input, maxUserLineLength),
		output:     output,
		cfg:        cfg,
		history:    history,
		serialized: serialized,

func (s *Session) Close() error {
	if s.serialized != nil {
		return s.serialized.Close()
	}
	return nil
}

func (s *Session) getUserInput(ctx context.Context) (string, error) {
	_, err := fmt.Fprintf(s.output, "> ")
	if err != nil {
		return "", err
	}
	input, err := s.input.ReadLine()
	if err != nil {
		if errors.Is(err, io.EOF) {
			_, _ = fmt.Fprint(s.output, "\r")
		}
		return "", err
	}
	if !strings.HasSuffix(input, "\n") {
		_, _ = fmt.Fprint(s.output, "\n\n")
	} else {
		_, _ = fmt.Fprint(s.output, "\n")
	}

	return input, nil
}

func (s *Session) callTool(ctx context.Context, call tools.Call) (rv string) {
	start := time.Now()
	defer func() {
		rv = fmt.Sprintf("[start: %s, duration: %s]\n",
			time.Now().Format("2006-01-02 15:04:05 MST"),
			time.Since(start)) + rv
	}()

	if call.Ref == nil {
		return fmt.Sprintf("error: unknown tool %q", call.Name)
	}
	result, err := call.Ref.Function(ctx, call)
	if err != nil {
		return fmt.Sprintf("error: %v", err)
	}
	return result
}

func (s *Session) buildContextMessage() string {
	cwd, err := os.Getwd()
	if err != nil {
		cwd = "(unknown)"
	}
	return fmt.Sprintf("Working directory: %s\n", cwd)
}

func (s *Session) Run(ctx context.Context) error {
	addToHistory := func(p ...prompt.Prompt) error {
		s.history = append(s.history, p...)
		if s.serialized != nil {
			return s.serialized.Append(p...)
		}
		return nil
	}

	// Inject execution context at the start of every run
	if err := addToHistory(prompt.AsUser(s.buildContextMessage())); err != nil {
		return err
	}

	for firstLoop := true; true; firstLoop = false {
		if !firstLoop {
			resp, err := s.gen.WithContext(ctx).Prompt(s.history...)
			if err != nil {
				return err
			}

			for _, text := range resp.Texts {
				if _, err := fmt.Fprintf(s.output, "%s\n\n", text); err != nil {
					return err
				}
				if err := addToHistory(prompt.AsAssistant(text)); err != nil {
					return err
				}
			}

			for _, call := range resp.Tools {
				_, err = fmt.Fprintf(s.output, "[%s %s]\n", call.Name, jsonUnescapeHTML.Replace(string(call.Argument)))
				if err != nil {
					return err
				}
				if err := addToHistory(
					prompt.AsToolCal

func formatContextLines(lines []string, targetLine, contextRadius int) string {
	start := targetLine - contextRadius
	if start < 1 {
		start = 1
	}
	end := targetLine + contextRadius
	if end > len(lines) {
		end = len(lines)
	}
	var sb strings.Builder
	for i := start; i <= end; i++ {
		hash := hashLineContent(lines[i-1])
		fmt.Fprintf(&sb, "  %d:%s|%s\n", i, hash, lines[i-1])
	}
	return sb.String()
}

func readFileLines(path string) ([]string, bool, error) {
	data, err := os.ReadFile(path)
	if err != nil {
		return nil, false, err
	}
	content := string(data)
	if content == "" {
		return nil, false, nil
	}
	endsWithNewline := strings.HasSuffix(content, "\n")
	lines := strings.Split(content, "\n")
	if endsWithNewline && len(lines) > 0 && lines[len(lines)-1] == "" {
		lines = lines[:len(lines)-1]
	}
	return lines, endsWithNewline, nil
}

func formatHashlines(lines []string, lineOffset int) string {
	var sb strings.Builder
	for i, line := range lines {
		lineNum := lineOffset + i + 1 // 1-indexed
		hash := hashLineContent(line)
		fmt.Fprintf(&sb, "%d:%s|%s\n", lineNum, hash, line)
	}
	return sb.String()
}

func parseHashlineRef(ref string) (lineNum int, hash string, err error) {
	parts := strings.SplitN(ref, ":", 2)
	if len(parts) != 2 {
		return 0, "", fmt.Errorf("invalid hashline reference %q: expected format 'line:hash'", ref)
	}
	lineNum, err = strconv.Atoi(parts[0])
	if err != nil {
		return 0, "", fmt.Errorf("invalid line number in %q: %v", ref, err)
	}
	if lineNum < 1 {
		return 0, "", fmt.Errorf("line number must be >= 1, got %d", lineNum)
	}
	return lineNum, parts[1], nil
}

func validateHashlineRef(lines []string, lineNum int, expectedHash string) error {
	if lineNum > len(lines) {
		return fmt.Errorf("line %d does not exist (file has %d lines)", lineNum, len(lines))
	}
	actualHash := hashLineContent(lines[lineNum-1])
	if actualHash != expectedHash {
		return fmt.Errorf("hash mismatch at line %d: expected %s, got %s (file has changed since last read)", lineNum, expectedHash, actualHash)
	}
	return nil
}

func formatDirEntry(path string, entry os.DirEntry) string {
	info, err := entry.Info()
	if err != nil {
		return fmt.Sprintf("  ?  %s", entry.Name())
	}
	return fmt.Sprintf("%s  %10d  %s  %s",
		info.Mode().String(),
		info.Size(),
		info.ModTime().Format("2006-01-02 15:04"),
		entry.Name(),
	)
}

func buildTree(lines *[]string, dir, prefix string, depth int, showHidden bool) {
	if depth <= 0 {
		return
	}

	entries, err := os.ReadDir(dir)
	if err != nil {
		return
	}

	// Filter and sort: directories first, then files, alphabetical within each group
	var filtered []treeEntry
	for _, e := range entries {
		name := e.Name()
		if !showHidden && strings.HasPrefix(name, ".") {
			continue
		}
		filtered = append(filtered, treeEntry{name: name, isDir: e.IsDir()})
	}
	sort.Slice(filtered, func(i, j int) bool {
		if filtered[i].isDir != filtered[j].isDir {
			return filtered[i].isDir
		}
		return filtered[i].name < filtered[j].name
	})

	for i, entry := range filtered {
		isLast := i == len(filtered)-1
		connector := "├── "
		childPrefix := "│   "
		if isLast {
			connector = "└── "
			childPrefix = "    "
		}

		*lines = append(*lines, prefix+connector+entry.name)

		if entry.isDir {
			buildTree(lines, filepath.Join(dir, entry.name), prefix+childPrefix, depth-1, showHidden)
		}
	}
}

func NewWebSearchTool(apiKey, endpoint string) tools.Tool {
	if endpoint == "" {
		endpoint = defaultSearchEndpoint
	}
	return tools.NewTool("web_search",
		tools.WithDescription("Search the web using Brave Search and return the top results."),
		tools.WithArgSchema(searchArgs{}),
		tools.WithFunction(func(ctx context.Context, call tools.Call) (string, error) {
			var params searchArgs
			if err := json.Unmarshal(call.Argument, &params); err != nil {
				return fmt.Sprintf("error: invalid arguments: %v", err), nil
			}
			if params.Query == "" {
				return "error: query is required", nil
			}

			reqURL := endpoint + "?q=" + url.QueryEscape(params.Query)
			req, err := http.NewRequestWithContext(ctx, http.MethodGet, reqURL, nil)
			if err != nil {
				return fmt.Sprintf("error: %v", err), nil
			}
			req.Header.Set("Accept", "application/json")
			req.Header.Set("X-Subscription-Token", apiKey)

			client := &http.Client{Timeout: fetchTimeout}
			resp, err := client.Do(req)
			if err !=

func defaultSessionPath() (string, error) {
	home, err := os.UserHomeDir()
	if err != nil {
		return "", fmt.Errorf("getting home directory: %w", err)
	}
	dir := filepath.Join(home, ".ajent", "sessions")
	if err := os.MkdirAll(dir, 0755); err != nil {
		return "", fmt.Errorf("creating sessions directory: %w", err)
	}

	cwd, err := os.Getwd()
	if err != nil {
		return "", fmt.Errorf("getting working directory: %w", err)
	}
	base := filepath.Base(cwd)

	var b [4]byte
	if _, err := rand.Read(b[:]); err != nil {
		return "", fmt.Errorf("generating random bytes: %w", err)
	}
	hash := hex.EncodeToString(b[:])

	ts := time.Now().Format("20060102-150405")
	name := fmt.Sprintf("%s-%s-%s.hjl", base, ts, hash)
	return filepath.Join(dir, name), nil
}