ollie/cmd/olliesrv/internal/agent/background.go

package agent

import (
	"fmt"
	"strings"
	"sync"

	"ollie/toolsrv"
)

// bgProc tracks a background process for interrupt injection.
type bgProc struct {
	PID        int
	Cmd        string // human-readable command description
	Tool       string // tool name that started it
	LastOutput string // last seen output (for diff detection)
}

// bgTracker manages background process state for an agent.
type bgTracker struct {
	mu    sync.Mutex
	procs []bgProc
}

// Add registers a new background process.
func (bt *bgTracker) Add(pid int, tool, cmd string) {
	bt.mu.Lock()
	defer bt.mu.Unlock()
	bt.procs = append(bt.procs, bgProc{PID: pid, Tool: tool, Cmd: cmd})
}

// Remove removes a process by PID (after it's been reported as exited).
func (bt *bgTracker) Remove(pid int) {
	bt.mu.Lock()
	defer bt.mu.Unlock()
	for i, p := range bt.procs {
		if p.PID == pid {
			bt.procs = append(bt.procs[:i], bt.procs[i+1:]...)
			return
		}
	}
}

// CollectInterrupts checks all tracked background procs for new output
// and returns formatted interrupt blocks. Removes completed procs after
// reporting their final output.
func (bt *bgTracker) CollectInterrupts(conn *toolsrv.Conn) string {
	if conn == nil {
		return ""
	}
	bt.mu.Lock()
	defer bt.mu.Unlock()

	if len(bt.procs) == 0 {
		return ""
	}

	var sb strings.Builder
	var completed []int

	for i := range bt.procs {
		p := &bt.procs[i]

		// Get current output
		output, err := conn.GetDetachedOutput(p.PID)
		if err != nil {
			// Process may have been dismissed or doesn't exist
			completed = append(completed, i)
			continue
		}

		// Only inject if output changed
		if output == p.LastOutput {
			continue
		}
		p.LastOutput = output

		// Get process status
		status := "running"
		exitCode := -1

		// Check if exited by trying to parse stat
		info := readProcStat(conn, p.PID)
		if info != nil {
			if exited, ok := info["exited"].(bool); ok && exited {
				status = "exited"
				if code, ok := info["exit_code"].(int); ok {
					exitCode = code
				}
				completed = append(completed, i)
			}
		}

		// Tail last 20 lines
		lines := tailLines(output, 20)

		sb.WriteString("\n\n<system-proc-interrupt id=\"")
		sb.WriteString(fmt.Sprintf("%d", p.PID))
		sb.WriteString("\" cmd=\"")
		sb.WriteString(escapeAttr(p.Cmd))
		sb.WriteString("\" status=\"")
		sb.WriteString(status)
		sb.WriteString("\"")
		if exitCode >= 0 {
			sb.WriteString(fmt.Sprintf(" exit=\"%d\"", exitCode))
		}
		sb.WriteString(">\n")
		sb.WriteString(lines)
		sb.WriteString("\n</system-proc-interrupt>")
	}

	// Remove completed procs (reverse order to preserve indices)
	for i := len(completed) - 1; i >= 0; i-- {
		idx := completed[i]
		bt.procs = append(bt.procs[:idx], bt.procs[idx+1:]...)
	}

	return sb.String()
}

// HasProcs returns true if there are any tracked background processes.
func (bt *bgTracker) HasProcs() bool {
	bt.mu.Lock()
	defer bt.mu.Unlock()
	return len(bt.procs) > 0
}

// readProcStat reads status info for a background proc.
func readProcStat(conn *toolsrv.Conn, pid int) map[string]any {
	// Use the ListDetachedRaw to find this proc
	procs := conn.ListDetachedRaw()
	for _, p := range procs {
		if m, ok := p.(map[string]any); ok {
			if mpid, ok := m["pid"].(int); ok && mpid == pid {
				return m
			}
		}
	}
	return nil
}

// tailLines returns the last n lines of s.
func tailLines(s string, n int) string {
	lines := strings.Split(strings.TrimRight(s, "\n"), "\n")
	if len(lines) > n {
		lines = lines[len(lines)-n:]
	}
	return strings.Join(lines, "\n")
}

// escapeAttr escapes a string for use in an XML-like attribute.
func escapeAttr(s string) string {
	s = strings.ReplaceAll(s, "&", "&amp;")
	s = strings.ReplaceAll(s, "\"", "&quot;")
	s = strings.ReplaceAll(s, "<", "&lt;")
	return s
}