orchestra/internal/herdr/occupancy.go

package herdr

import (
	"bufio"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"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
}
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
}

// 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") }