template-jj/cmd/planctl/context.go

package main

import (
	"bufio"
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"strconv"
	"strings"
)

// spec:26174-planctl-context-tokens/R2.3+D§3.5
var errNoUsageLine = errors.New("no assistant line with usage object")

// spec:26174-planctl-context-tokens/D§3.3
// pidSessionDir is the on-disk root for PID-keyed session files, written by the
// SessionStart hook. Injected as a variable so tests can stage fixtures in t.TempDir().
var pidSessionDir = "/tmp/planctl-sessions"

// pidWalkMaxHops is the ceiling on process-ancestor lookups during PID-file search (R1.2).
const pidWalkMaxHops = 8

// spec:26174-planctl-context-tokens/R1.1+R1.2+R1.3+R1.4+R1.5+R1.6+R2.1+R2.2+R2.3+R2.4+R2.5+R2.6+D§3.2+D§6.4
// ReadTokenCtx is the ambient constructor: it reads the calling process's
// environment and/or ancestor PID files to discover the session UUID, then
// opens the transcript and extracts the most recent API usage. Returns nil
// when any step fails — callers treat nil as "no context data". Never invokes
// a subprocess (design §6.4).
func ReadTokenCtx() *TokenCtx {
	uuid := resolveSessionUUID()
	if uuid == "" {
		return nil
	}
	home, _ := os.UserHomeDir()
	if home == "" {
		return nil
	}
	path := transcriptPath(uuid, home, os.Getenv("CLAUDE_PROJECT_DIR"))
	used, err := parseLastUsage(path)
	if err != nil {
		return nil
	}
	return &TokenCtx{Used: used, Limit: readLimit()}
}

// spec:26174-planctl-context-tokens/R1.1+R1.2+R1.3+R1.4+R6.1+R6.3+D§3.3
// resolveSessionUUID returns the session UUID string, or "" if not found.
// Priority: CLAUDE_SESSION_ID env > PID-ancestor walk with liveness + start-stamp checks.
func resolveSessionUUID() string {
	if v := os.Getenv("CLAUDE_SESSION_ID"); v != "" {
		return v
	}
	return resolveViaPIDWalk(os.Getpid(), pidSessionDir, ppidOf, processAlive, processStartTimeSec)
}

// resolveViaPIDWalk does the ancestor-PID walk with pluggable platform helpers
// so tests can feed a synthetic chain without needing real processes.
func resolveViaPIDWalk(
	startPID int,
	dir string,
	parentOf func(int) int,
	alive func(int) bool,
	startStampOf func(int) (int64, error),
) string {
	pid := startPID
	for hop := 0; hop < pidWalkMaxHops; hop++ {
		pid = parentOf(pid)
		if pid <= 1 {
			return ""
		}
		path := filepath.Join(dir, strconv.Itoa(pid))
		uuid, stamp, err := parsePIDFile(path)
		if err != nil {
			continue
		}
		if !alive(pid) {
			continue
		}
		actual, err := startStampOf(pid)
		if err != nil || actual != stamp {
			continue
		}
		return uuid
	}
	return ""
}

// spec:26174-planctl-context-tokens/D§4.2
type transcriptUsage struct {
	Type    string `json:"type"`
	Message struct {
		Usage *struct {
			InputTokens              int64 `json:"input_tokens"`
			CacheCreationInputTokens int64 `json:"cache_creation_input_tokens"`
			CacheReadInputTokens     int64 `json:"cache_read_input_tokens"`
		} `json:"usage"`
	} `json:"message"`
}

// spec:26174-planctl-context-tokens/R2.1+R2.2+R2.3+R2.6+D§3.5
// parseLastUsage scans a JSONL transcript and returns the summed token count
// (input + cache_read + cache_creation) from the *last* line whose top-level
// "type" is "assistant" and whose "message.usage" object is present.
// Presence qualifies (per PRD R2.1) — the usage object may contain zeros.
// Malformed JSON lines are silently skipped. Returns errNoUsageLine when no
// qualifying line exists.
func parseLastUsage(path string) (int64, error) {
	f, err := os.Open(path)
	if err != nil {
		return 0, err
	}
	defer f.Close()
	sc := bufio.NewScanner(f)
	sc.Buffer(make([]byte, 64*1024), 10<<20)
	var last int64
	var found bool
	for sc.Scan() {
		var rec transcriptUsage
		if err := json.Unmarshal(sc.Bytes(), &rec); err != nil {
			continue
		}
		if rec.Type != "assistant" || rec.Message.Usage == nil {
			continue
		}
		u := rec.Message.Usage
		last = u.InputTokens + u.CacheCreationInputTokens + u.CacheReadInputTokens
		found = true
	}
	if err := sc.Err(); err != nil {
		return 0, err
	}
	if !found {
		return 0, errNoUsageLine
	}
	return last, nil
}

// spec:26174-planctl-context-tokens/R6.2+R6.3+D§3.9+D§4.1
// parsePIDFile reads a PID file matching the exact layout in design §4.1:
//
//	<session-uuid>\n
//	<start-time-stamp>\n
//
// Start-time stamp is platform-specific (clock ticks since boot on Linux, epoch
// seconds on macOS — D§4.1 / §5.5). The parser is strict: both lines must be
// newline-terminated, uuid must be non-empty, stamp must parse as int64, and
// nothing may follow the second newline. Anything else (partial write, extra
// lines, missing final \n) is rejected — a corrupted file should not pass the
// liveness/start-stamp check in resolveViaPIDWalk.
func parsePIDFile(path string) (uuid string, startStamp int64, err error) {
	data, err := os.ReadFile(path)
	if err != nil {
		return "", 0, err
	}
	rest, ok := strings.CutSuffix(string(data), "\n")
	if !ok {
		return "", 0, fmt.Errorf("pid file %s: missing terminating newline", path)
	}
	uuid, stampPart, ok := strings.Cut(rest, "\n")
	if !ok {
		return "", 0, fmt.Errorf("pid file %s: missing line 2 (start stamp)", path)
	}
	if uuid == "" {
		return "", 0, fmt.Errorf("pid file %s: empty uuid", path)
	}
	if strings.Contains(stampPart, "\n") {
		return "", 0, fmt.Errorf("pid file %s: unexpected trailing content", path)
	}
	stamp, perr := strconv.ParseInt(stampPart, 10, 64)
	if perr != nil {
		return "", 0, fmt.Errorf("pid file %s: malformed start stamp: %w", path, perr)
	}
	return uuid, stamp, nil
}

// spec:26174-planctl-context-tokens/R1.5+R1.6+D§3.4
// transcriptPath assembles the Claude Code transcript path.
// slug = projectDir with every '/' replaced by '-'. Falls back to os.Getwd() when projectDir is "".
func transcriptPath(uuid, home, projectDir string) string {
	if projectDir == "" {
		projectDir, _ = os.Getwd()
	}
	slug := strings.ReplaceAll(projectDir, "/", "-")
	return filepath.Join(home, ".claude", "projects", slug, uuid+".jsonl")
}

// spec:26174-planctl-context-tokens/R2.4+D§3.6
// readLimit returns CLAUDE_CODE_MAX_CONTEXT_TOKENS as an int64, or 0 if unset/invalid.
func readLimit() int64 {
	v := os.Getenv("CLAUDE_CODE_MAX_CONTEXT_TOKENS")
	if v == "" {
		return 0
	}
	n, err := strconv.ParseInt(v, 10, 64)
	if err != nil || n < 0 {
		return 0
	}
	return n
}

// spec:26174-planctl-context-tokens/R2.2+D§3.1
// TokenCtx holds the resolved token-usage data for the current session.
// Limit == 0 means the limit is unknown (CLAUDE_CODE_MAX_CONTEXT_TOKENS absent).
// A nil *TokenCtx means no session data was available; all emitters must
// handle nil gracefully (emit nothing extra).
type TokenCtx struct {
	Used  int64
	Limit int64
}

// spec:26174-planctl-context-tokens/R4.4+D§3.8
func formatTokenCount(n int64) string {
	if n < 1000 {
		return "< 1 k"
	}
	k := (n + 500) / 1000
	return strconv.FormatInt(k, 10) + " k"
}

// spec:26174-planctl-context-tokens/R3.1+R3.2+R3.3+R3.5+D§3.7
// Classify returns the band's severity label and recommendation suffix.
// Returns ("", "") for nil receiver, unknown limit (Limit == 0), or Normal band (pct < 70).
func (t *TokenCtx) Classify() (severity, msg string) {
	if t == nil || t.Limit == 0 {
		return "", ""
	}
	pct := t.Used * 100 / t.Limit
	switch {
	case pct < 70:
		return "", ""
	case pct < 85:
		return "info", "plan to wrap up this session soon."
	case pct < 95:
		return "warn", "commit current work and start a new session after this task."
	default:
		return "error", "stop new work; commit and close out immediately."
	}
}