ollie/cmd/Ollie/main.go

// ollie-acme: acme interface for ollie AI sessions.
// All file operations use 9P directly (not FUSE mount) for proper
// blocking read support on streaming files like chat.
package main

import (
	"io"
	"log"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"sync"

	"9fans.net/go/acme"
	"9fans.net/go/plan9"
	"9fans.net/go/plan9/client"
)

func main() {
	acme.AutoExit(true)
	if err := newSessionList(); err != nil {
		log.Fatal(err)
	}
	select {}
}

// --- 9P helpers ---

func dial9p() *client.Fsys {
	ns := os.Getenv("NAMESPACE")
	if ns == "" {
		ns = client.Namespace()
	}
	if ns == "" {
		log.Fatal("cannot determine namespace; set $NAMESPACE")
	}
	address := ns + "/ollie"

	conn, err := client.Dial("unix", address)
	if err != nil {
		log.Fatalf("dial %s: %v", address, err)
	}
	fsys, err := conn.Attach(nil, getuser(), "")
	if err != nil {
		log.Fatalf("attach: %v", err)
	}
	return fsys
}

func getuser() string {
	if u := os.Getenv("USER"); u != "" {
		return u
	}
	return "none"
}

func read9p(path string) ([]byte, error) {
	fsys := dial9p()
	fid, err := fsys.Open(path, plan9.OREAD)
	if err != nil {
		return nil, err
	}
	defer fid.Close()
	var buf []byte
	tmp := make([]byte, 8192)
	for {
		n, err := fid.Read(tmp)
		if n > 0 {
			buf = append(buf, tmp[:n]...)
		}
		if err != nil {
			if err == io.EOF {
				break
			}
			return buf, err
		}
	}
	return buf, nil
}

func write9p(path string, data []byte) error {
	fsys := dial9p()
	fid, err := fsys.Open(path, plan9.OWRITE)
	if err != nil {
		return err
	}
	defer fid.Close()
	_, err = fid.Write(data)
	return err
}

// agentPath returns the 9P path to the active agent within a session.
func agentPath(sid string) string {
	fsys := dial9p()
	dirPath := "session/" + sid + "/agent"
	fid, err := fsys.Open(dirPath, plan9.OREAD)
	if err != nil {
		return dirPath + "/0"
	}
	dirs, err := fid.Dirreadall()
	fid.Close()
	if err != nil || len(dirs) == 0 {
		return dirPath + "/0"
	}
	for _, d := range dirs {
		if d.Mode&plan9.DMDIR != 0 {
			return dirPath + "/" + d.Name
		}
	}
	return dirPath + "/" + dirs[0].Name
}

// --- session list window ---

type sessionList struct {
	win     *acme.Win
	allSIDs []string
	mu      sync.Mutex
}

func newSessionList() error {
	w, err := acme.New()
	if err != nil {
		return err
	}
	w.Name("ollie/sessions")
	sl := &sessionList{win: w}
	sl.refresh()
	w.Ctl("cleartag")
	w.Write("tag", []byte(" New Kill Refresh"))
	go sl.eventLoop()
	return nil
}

func (sl *sessionList) refresh() {
	data, err := read9p("session/idx")
	if err != nil {
		sl.win.Errf("read idx: %v", err)
		return
	}
	sl.win.Addr(",")
	sl.win.Write("data", nil)

	raw := strings.TrimSpace(string(data))
	if raw == "" {
		sl.win.Ctl("clean")
		return
	}

	lines := strings.Split(raw, "\n")
	sort.Strings(lines)

	sl.mu.Lock()
	sl.allSIDs = nil
	for _, l := range lines {
		if s := strings.TrimSpace(l); s != "" {
			id, _, _ := strings.Cut(s, "\t")
			sl.allSIDs = append(sl.allSIDs, id)
		}
	}
	sl.mu.Unlock()

	sl.win.Write("body", []byte(raw+"\n"))
	sl.win.Ctl("clean")
}

func (sl *sessionList) resolveFullSID(text string) string {
	text = strings.TrimSpace(text)
	if text == "" {
		return ""
	}
	sl.mu.Lock()
	defer sl.mu.Unlock()
	for _, sid := range sl.allSIDs {
		parts := strings.Split(sid, "__")
		if parts[len(parts)-1] == text {
			return sid
		}
	}
	return ""
}

func (sl *sessionList) eventLoop() {
	for e := range sl.win.EventChan() {
		switch e.C2 {
		case 'x', 'X':
			cmd := strings.TrimSpace(string(e.Text))
			if cmd == "New" {
				openNewSession()
				continue
			}
			if cmd == "Kill" {
				arg := strings.TrimSpace(string(e.Arg))
				sid := sl.resolveFullSID(arg)
				if sid != "" {
					ab := agentPath(sid)
					write9p(ab+"/ctl", []byte("kill"))
					sl.refresh()
				}
				continue
			}
			if cmd == "Refresh" {
				sl.refresh()
				continue
			}
			sid := sl.resolveFullSID(cmd)
			if sid != "" {
				openChat(sid)
			} else {
				sl.win.WriteEvent(e)
			}
		case 'l', 'L':
			sl.win.WriteEvent(e)
		}
	}
	os.Exit(0)
}

// --- chat window ---

type chatWin struct {
	win *acme.Win
	sid string
	mu  sync.Mutex
}

var chatWindows sync.Map

func openChat(sid string) {
	if v, ok := chatWindows.Load(sid); ok {
		v.(*chatWin).win.Ctl("show")
		return
	}

	w, err := acme.New()
	if err != nil {
		log.Printf("open chat %s: %v", sid, err)
		return
	}
	winName := "ollie/sessions/" + sid
	if data, err := read9p("session/idx"); err == nil {
		for _, line := range strings.Split(string(data), "\n") {
			fields := strings.Split(line, "\t")
			if len(fields) >= 3 && fields[0] == sid {
				winName = fields[2] + "/+Ollie"
				break
			}
		}
	}
	w.Name("%s", winName)
	w.Ctl("cleartag")
	w.Fprintf("tag", " Prompt Stop Ctl Cfg Nudge Usage")

	cw := &chatWin{win: w, sid: sid}
	chatWindows.Store(sid, cw)

	// Load existing log (last 64KB) for initial context
	ab := agentPath(sid)
	if data, err := read9p(ab + "/log"); err == nil && len(data) > 0 {
		w.Write("body", data)
	}
	w.Ctl("clean")

	go cw.streamChat(ab)
	go cw.eventLoop(ab)
}

func (cw *chatWin) eventLoop(ab string) {
	defer chatWindows.Delete(cw.sid)
	for e := range cw.win.EventChan() {
		switch e.C2 {
		case 'x', 'X':
			cmd := strings.TrimSpace(string(e.Text))
			switch cmd {
			case "Prompt":
				openPrompt(ab)
			case "Stop":
				write9p(ab+"/ctl", []byte("stop"))
			case "Ctl":
				openFile9p(ab + "/ctl")
			case "Cfg":
				openFile9p(ab + "/cfg")
			case "Nudge":
				write9p(ab+"/prompt", []byte("."))
			case "Usage":
				if data, err := read9p(ab + "/usage"); err == nil {
					cw.win.Errf("%s", string(data))
				}
			default:
				cw.win.WriteEvent(e)
			}
		case 'l', 'L':
			cw.win.WriteEvent(e)
		}
	}
}

// streamChat reads the chat file via 9P (blocking reads) and appends to the window.
// The read blocks until new tokens arrive, then returns them. Loops until EOF
// (session killed) or error.
func (cw *chatWin) streamChat(ab string) {
	fsys := dial9p()
	fid, err := fsys.Open(ab+"/chat", plan9.OREAD)
	if err != nil {
		log.Printf("open chat stream: %v", err)
		return
	}
	defer fid.Close()

	buf := make([]byte, 8192)
	for {
		n, err := fid.Read(buf)
		if n > 0 {
			cw.mu.Lock()
			cw.win.Addr("$")
			cw.win.Write("data", buf[:n])
			cw.win.Ctl("clean")
			cw.mu.Unlock()
		}
		if err != nil {
			// EOF = session killed; any other error = done
			return
		}
	}
}

// --- helper windows ---

func openNewSession() {
	w, err := acme.New()
	if err != nil {
		return
	}
	w.Name("ollie/new")
	w.Write("body", []byte("cwd="+os.Getenv("HOME")+"\n"))
	w.Ctl("cleartag")
	w.Fprintf("tag", " Put")
	go func() {
		for e := range w.EventChan() {
			if e.C2 == 'x' || e.C2 == 'X' {
				cmd := strings.TrimSpace(string(e.Text))
				if cmd == "Put" {
					w.Addr(",")
					data, _ := w.ReadAll("body")
					write9p("session/new", data)
					w.Del(true)
					return
				}
			}
		}
	}()
}

func openPrompt(ab string) {
	w, err := acme.New()
	if err != nil {
		return
	}
	w.Name("%s", filepath.Base(ab)+"/prompt")
	w.Ctl("cleartag")
	w.Fprintf("tag", " Put")
	go func() {
		for e := range w.EventChan() {
			if e.C2 == 'x' || e.C2 == 'X' {
				cmd := strings.TrimSpace(string(e.Text))
				if cmd == "Put" {
					w.Addr(",")
					data, _ := w.ReadAll("body")
					write9p(ab+"/prompt", data)
					w.Addr(",")
					w.Write("data", nil)
					w.Ctl("clean")
				}
			}
		}
	}()
}

func openFile9p(path string) {
	w, err := acme.New()
	if err != nil {
		return
	}
	parts := strings.Split(path, "/")
	w.Name("%s", parts[len(parts)-1])
	if data, err := read9p(path); err == nil {
		w.Write("body", data)
	}
	w.Ctl("cleartag")
	w.Fprintf("tag", " Put Get")
	go func() {
		for e := range w.EventChan() {
			if e.C2 == 'x' || e.C2 == 'X' {
				cmd := strings.TrimSpace(string(e.Text))
				switch cmd {
				case "Put":
					w.Addr(",")
					data, _ := w.ReadAll("body")
					write9p(path, data)
					w.Ctl("clean")
				case "Get":
					if data, err := read9p(path); err == nil {
						w.Addr(",")
						w.Write("data", data)
						w.Ctl("clean")
					}
				}
			}
		}
	}()
}