orchestra/internal/herdr/occupancy.go

package herdr

import (
	"bufio"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"strconv"
	"strings"
	"time"
)

type Usage struct{ Input, CacheRead, CacheWrite, Output int64 }

func (u Usage) Numerator() int64 { return u.Input + u.CacheRead + u.CacheWrite }
func Fraction(u Usage, w int64) float64 {
	if w <= 0 {
		return 0
	}
	f := float64(u.Numerator()) / float64(w)
	if f < 0 {
		return 0
	}
	if f > 1 {
		return 1
	}
	return f
}

// ClaudeSessionFile resolves the transcript file for a Claude Code session
// running against worktree, by newest-mtime under Claude Code's encoded
// project directory (~/.claude/projects/<abs-worktree-path-with-/-replaced-
// by-minus>/*.jsonl). This is the "resolve at lease time" fallback called
// out in the spec (§5.2.1); the Stop hook's transcript_path (see
// ClaudeStopHookUsage) is the authoritative source once wired.
func ClaudeSessionFile(worktree string) (string, error) {
	abs, err := filepath.Abs(worktree)
	if err != nil {
		return "", err
	}
	enc := strings.ReplaceAll(abs, "/", "-")
	home := os.Getenv("HOME")
	dir := filepath.Join(home, ".claude", "projects", enc)
	matches, err := filepath.Glob(filepath.Join(dir, "*.jsonl"))
	if err != nil {
		return "", err
	}
	if len(matches) == 0 {
		return "", fmt.Errorf("claude: no session transcripts found in %s", dir)
	}
	var newest string
	var newestMod time.Time
	for _, m := range matches {
		fi, err := os.Stat(m)
		if err != nil {
			continue
		}
		if fi.ModTime().After(newestMod) {
			newestMod = fi.ModTime()
			newest = m
		}
	}
	if newest == "" {
		return "", fmt.Errorf("claude: could not stat any session transcript in %s", dir)
	}
	return newest, nil
}

func ClaudeUsage(p string) (Usage, error) {
	f, e := os.Open(p)
	if e != nil {
		return Usage{}, e
	}
	defer f.Close()
	s := bufio.NewScanner(f)
	var last Usage
	for s.Scan() {
		var x struct {
			Message struct {
				Usage struct {
					Input  int64 `json:"input_tokens"`
					Read   int64 `json:"cache_read_input_tokens"`
					Write  int64 `json:"cache_creation_input_tokens"`
					Output int64 `json:"output_tokens"`
				} `json:"usage"`
			} `json:"message"`
		}
		if json.Unmarshal(s.Bytes(), &x) == nil && x.Message.Usage.Input > 0 {
			last = Usage{x.Message.Usage.Input, x.Message.Usage.Read, x.Message.Usage.Write, x.Message.Usage.Output}
		}
	}
	return last, s.Err()
}

// ClaudeStopHook is the JSON contract supplied by Claude Code's stop hook.
// Keeping the hook parser here makes ingestion independent of the hook's shell.
type ClaudeStopHook struct {
	TranscriptPath string `json:"transcript_path"`
}

func ClaudeStopHookUsage(r io.Reader) (Usage, string, error) {
	var h ClaudeStopHook
	if err := json.NewDecoder(r).Decode(&h); err != nil {
		return Usage{}, "", err
	}
	if h.TranscriptPath == "" {
		return Usage{}, "", fmt.Errorf("claude stop hook: transcript_path required")
	}
	u, err := ClaudeUsage(h.TranscriptPath)
	return u, h.TranscriptPath, err
}

// CodexRolloutPaths discovers active rollouts from configured Codex state.
// sqlite3 is intentionally used as an optional bridge: Codex owns the schema
// and deployments may not ship a Go sqlite driver.
func CodexRolloutPaths(home string) ([]string, error) {
	if home == "" {
		home = os.Getenv("CODEX_HOME")
	}
	if home == "" {
		home = filepath.Join(os.Getenv("HOME"), ".codex")
	}
	var paths []string
	matches, _ := filepath.Glob(filepath.Join(home, "state_*.sqlite"))
	for _, db := range matches {
		out, err := exec.Command("sqlite3", db, "select rollout_path from threads where rollout_path is not null;").Output()
		if err == nil {
			for _, p := range strings.Fields(string(out)) {
				if p != "" {
					paths = append(paths, p)
				}
			}
		}
	}
	if len(paths) == 0 {
		paths, _ = filepath.Glob(filepath.Join(home, "sessions", "*", "*", "*", "rollout-*.jsonl"))
	}
	if len(paths) == 0 {
		return nil, fmt.Errorf("codex: no active rollout found in %s", home)
	}
	return paths, nil
}
func CodexActiveUsage(home string) (Usage, string, error) {
	paths, err := CodexRolloutPaths(home)
	if err != nil {
		return Usage{}, "", err
	}
	for i := len(paths) - 1; i >= 0; i-- {
		if _, e := os.Stat(paths[i]); e == nil {
			u, e := CodexUsage(paths[i])
			return u, paths[i], e
		}
	}
	return Usage{}, "", os.ErrNotExist
}
func CodexUsage(p string) (Usage, error) {
	f, e := os.Open(p)
	if e != nil {
		return Usage{}, e
	}
	defer f.Close()
	s := bufio.NewScanner(f)
	var u Usage
	for s.Scan() {
		var x struct {
			Payload struct {
				Type string `json:"type"`
				Info struct {
					Last struct {
						Input int64 `json:"input"`
						Read  int64 `json:"cached_input"`
					} `json:"last_token_usage"`
				} `json:"info"`
			} `json:"payload"`
		}
		if json.Unmarshal(s.Bytes(), &x) == nil && x.Payload.Type == "token_count" {
			u = Usage{x.Payload.Info.Last.Input, x.Payload.Info.Last.Read, 0, 0}
		}
	}
	return u, s.Err()
}
func OpenCodeUsage(p string) (Usage, error) {
	f, e := os.Open(p)
	if e != nil {
		return Usage{}, e
	}
	defer f.Close()
	var x struct {
		Tokens struct {
			Input  int64 `json:"input"`
			Output int64 `json:"output"`
			Cache  struct {
				Read  int64 `json:"read"`
				Write int64 `json:"write"`
			} `json:"cache"`
		} `json:"tokens"`
	}
	e = json.NewDecoder(f).Decode(&x)
	return Usage{x.Tokens.Input, x.Tokens.Cache.Read, x.Tokens.Cache.Write, x.Tokens.Output}, e
}

// OpenCodeSessionID resolves the exact OpenCode session associated with a
// checkout. OpenCode stores usage in its SQLite session table, not in a pane
// transcript. Selecting by directory and persisting the returned id prevents
// a multi-pane worker from attributing another task's newest session to this
// lease.
func OpenCodeSessionID(worktree string) (string, error) {
	db := os.Getenv("ORCHESTRA_OPENCODE_DB")
	if db == "" {
		db = filepath.Join(os.Getenv("HOME"), ".local", "share", "opencode", "opencode.db")
	}
	abs, err := filepath.Abs(worktree)
	if err != nil {
		return "", err
	}
	out, err := exec.Command("sqlite3", "-readonly", "-noheader", db, "select id from session where directory = "+sqliteQuote(abs)+" order by time_updated desc limit 1;").Output()
	if err != nil {
		return "", fmt.Errorf("opencode session lookup: %w", err)
	}
	id := strings.TrimSpace(string(out))
	if id == "" {
		return "", fmt.Errorf("opencode session lookup: no session for %s", abs)
	}
	return id, nil
}

// OpenCodeSessionUsage reads the token counters for one persisted session.
func OpenCodeSessionUsage(sessionID string) (Usage, error) {
	if sessionID == "" {
		return Usage{}, fmt.Errorf("opencode session id required")
	}
	db := os.Getenv("ORCHESTRA_OPENCODE_DB")
	if db == "" {
		db = filepath.Join(os.Getenv("HOME"), ".local", "share", "opencode", "opencode.db")
	}
	out, err := exec.Command("sqlite3", "-readonly", "-noheader", "-separator", "|", db, "select tokens_input,tokens_cache_read,tokens_cache_write,tokens_output from session where id = "+sqliteQuote(sessionID)+";").Output()
	if err != nil {
		return Usage{}, fmt.Errorf("opencode usage lookup: %w", err)
	}
	parts := strings.Split(strings.TrimSpace(string(out)), "|")
	if len(parts) != 4 {
		return Usage{}, fmt.Errorf("opencode usage lookup: unknown session %q", sessionID)
	}
	values := [4]int64{}
	for i, part := range parts {
		values[i], err = strconv.ParseInt(part, 10, 64)
		if err != nil {
			return Usage{}, fmt.Errorf("opencode usage lookup: %w", err)
		}
	}
	return Usage{Input: values[0], CacheRead: values[1], CacheWrite: values[2], Output: values[3]}, nil
}

func sqliteQuote(s string) string { return "'" + strings.ReplaceAll(s, "'", "''") + "'" }

// OpenCodeStatus probes the server fast path. Callers can use the returned
// status and fall back to OpenCodeUsage when the SSE/server is unavailable.
func OpenCodeStatus(ctx context.Context, baseURL, sessionID string) (string, error) {
	if baseURL == "" {
		baseURL = "http://127.0.0.1:4096"
	}
	u := strings.TrimRight(baseURL, "/") + "/session/" + sessionID
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, u, nil)
	if err != nil {
		return "", err
	}
	c := &http.Client{Timeout: 5 * time.Second}
	resp, err := c.Do(req)
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()
	if resp.StatusCode/100 != 2 {
		return "", fmt.Errorf("opencode status: http %s", resp.Status)
	}
	var x struct {
		Status string `json:"status"`
	}
	if err := json.NewDecoder(resp.Body).Decode(&x); err != nil {
		return "", err
	}
	return x.Status, nil
}
func IsBusy(s string) bool { return strings.EqualFold(s, "busy") }