diff --git a/AGENTS.md b/AGENTS.md index 6db5c98..bb21cc3 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -33,6 +33,14 @@ Instructions for AI agents working on this repository. - **internal/**: All implementation packages - **poc/**: Original proof-of-concept (reference only, not used in production) +## Adding Harness Support + +When adding support for a new AI agent tool (harness), read **`docs/adding-a-harness.md`** +before writing any code. It explains the key concepts, the expected file structure, and +important edge cases (thin hook payloads, async metrics, provisional capture). The existing +copilot harness (`internal/copilot/`, `internal/cli/app_copilot.go`) is the canonical +reference implementation to follow. + ## Testing - Unit tests: `go test ./...` diff --git a/docs/adding-a-harness.md b/docs/adding-a-harness.md new file mode 100644 index 0000000..4647611 --- /dev/null +++ b/docs/adding-a-harness.md @@ -0,0 +1,131 @@ +# Adding a New Harness to arok + +A **harness** is an AI agent tool that fires hooks at the end of a session — for example, +GitHub Copilot CLI, VS Code Copilot, Hermes, or OpenCode. This guide explains the concepts +you need to understand and the files you need to create. For concrete code structure, refer +to the existing harness implementations as living examples. + +## How capture works + +When an agent session ends, the harness fires a hook that runs: + +``` +arok capture --harness --event +``` + +arok receives a JSON payload from the hook, parses it, reads any additional session data +from the tool's own files or APIs, assembles a `session.SessionSummary`, and writes it to +the SQLite database. + +Your job as a harness author is to implement that payload→summary translation for your tool. + +## File structure + +Each harness lives in two places: + +``` +internal// ← payload parsing and session summarizing logic +internal/cli/app_.go ← CLI plumbing: runCapture() and helpers +``` + +**Read the copilot harness before writing any code.** It is the canonical, fully-tested +reference implementation: + +``` +internal/copilot/copilot.go ← Summarize(), ParsePayload(), ResolveSessionFile() +internal/cli/app_copilot.go ← runCaptureCopilot(), runReconcile(), and helpers +internal/cli/app_copilot_test.go ← integration tests via App.Run() +internal/copilot/testdata/ ← fixture JSONL files used by unit tests +``` + +The VS Code harness (`internal/vscode/`, `internal/cli/app_vscode.go`) is a simpler +example without reconcile logic — useful if your harness delivers complete data at hook time. + +## Key concepts + +### The hook payload may not contain everything you need + +Most harnesses pass only a thin payload (session ID, working directory, maybe an event +name). The real metrics — token counts, model names, tool calls — often live in a separate +log file or API that the tool writes during the session. + +Use the session ID and working directory from the payload as keys to locate and read that +richer data source. See how `ResolveSessionFile` in `internal/copilot/copilot.go` uses the +`COPILOT_HOME` environment variable and the session ID to find the right `events.jsonl` +file. + +### Final metrics may not be available when the hook fires + +Some tools emit usage metrics asynchronously — the hook fires to signal session end, but the +final token count is written to disk only moments later (or is computed by a background +process). If you try to read the metrics immediately, you will get incomplete data. + +**Do not block the hook.** The hook runner expects a fast exit. Instead: + +1. Capture whatever is available immediately and store it with + `CaptureState: "provisional"`. +2. Spawn a short-lived background process (`arok reconcile --harness `) that polls + until the final data appears, then upgrades the record to `CaptureState: "final"`. +3. If the background process exhausts its retries without finding complete data (e.g. the + process was killed), mark the record `CaptureState: "best_effort"` so the data is still + usable but clearly flagged as incomplete. + +This is exactly what the copilot harness does — see `spawnDetachedReconcile` and +`runReconcile` in `internal/cli/app_copilot.go`. + +If your tool delivers complete metrics synchronously at hook time, you can skip all of this +and always return `CaptureState: "final"`. The VS Code harness is an example of this +simpler path. + +### Use git metadata to enrich sessions + +The hook payload rarely includes repository context. Use `gitmeta.Inspect(cwd)` (see +`internal/gitmeta/`) to populate `RepoRoot`, `RepoBranch`, and `RepoRemote` on the +`SessionSummary` from the working directory that the payload provides. + +## Adding your harness: the steps + +1. **Create `internal//`** — implement `Summarize()`, which takes the parsed + payload and returns a `session.SessionSummary`. Model it on + `internal/copilot/copilot.go`. + +2. **Create `internal/cli/app_.go`** — implement `runCapture()` as a + method on `*App`. It reads the payload, calls `Summarize()`, opens the store, and calls + `db.UpsertSession()`. Model it on `internal/cli/app_copilot.go` (or the simpler + `app_vscode.go` if you don't need reconcile). + +3. **Wire the dispatcher** — add a `case "":` to the `switch` in `runCapture()` + in `internal/cli/app.go`. That single line is the only change needed in shared code. + +4. **Update the usage string** — add your harness name to the `--harness` list in + `printRootUsage()` in `app.go`. + +5. **Add `arok install` support (optional)** — if your tool uses a config file that arok + can write, follow the pattern in `runInstallCopilot` / `internal/install/copilot.go`. + +## Writing tests + +**Unit tests** — `internal//_test.go` + +Test `Summarize()` directly using fixture files in `internal//testdata/`. Your +fixture files are also living documentation of what the raw session data looks like — make +them representative. See `internal/copilot/copilot_test.go` and its testdata for the +pattern. + +**Integration tests** — `internal/cli/app__test.go` + +Test the end-to-end path through `App.Run()`. Every code path should be covered: a final +capture, a provisional capture (if applicable), and the failure cases (missing session ID, +unknown event, etc.). See `internal/cli/app_copilot_test.go` for the full pattern. + +Run `make check` before opening a PR. All tests must pass. + +## Checklist + +- [ ] `internal//.go` — `Summarize()` returns a valid `SessionSummary` +- [ ] `internal/cli/app_.go` — `runCapture()` method on `*App` +- [ ] `internal/cli/app.go` — one `case` added to the `runCapture()` switch +- [ ] `internal//_test.go` — unit tests for `Summarize()` +- [ ] `internal//testdata/` — fixture files covering final and (if needed) provisional states +- [ ] `internal/cli/app__test.go` — integration tests via `App.Run()` +- [ ] `make check` passes diff --git a/internal/cli/app.go b/internal/cli/app.go index de9550d..0bc3251 100644 --- a/internal/cli/app.go +++ b/internal/cli/app.go @@ -8,17 +8,12 @@ import ( "fmt" "io" "os" - "os/exec" "path/filepath" - "slices" "strings" - "syscall" "text/tabwriter" "time" "github.com/srbouffard/arok/internal/config" - "github.com/srbouffard/arok/internal/copilot" - "github.com/srbouffard/arok/internal/gitmeta" "github.com/srbouffard/arok/internal/install" sessionpkg "github.com/srbouffard/arok/internal/session" "github.com/srbouffard/arok/internal/store" @@ -81,62 +76,14 @@ func (a *App) runInstall(args []string) error { } } -func (a *App) runInstallCopilot(args []string) error { - fs := flag.NewFlagSet("install copilot", flag.ContinueOnError) - fs.SetOutput(a.stderr) - var ( - stateDirOverride string - copilotHome = install.DefaultCopilotHome() - binaryPath string - printConfig bool - ) - fs.StringVar(&stateDirOverride, "state-dir", "", "Override the AROK state directory.") - fs.StringVar(&copilotHome, "copilot-home", copilotHome, "Override the Copilot home directory.") - fs.StringVar(&binaryPath, "binary-path", "", "Override the binary path written into the Copilot hook config.") - fs.BoolVar(&printConfig, "print-config", false, "Print the generated Copilot hook config instead of installing it.") - if err := fs.Parse(args); err != nil { - return err - } - - stateDir, err := config.ResolveStateDir(stateDirOverride) - if err != nil { - return err - } - if err := config.EnsureLayout(stateDir); err != nil { - return err - } - - if binaryPath == "" { - binaryPath, err = os.Executable() - if err != nil { - return fmt.Errorf("resolve executable path: %w", err) - } - } - - db, err := store.Open(stateDir) - if err != nil { - return err - } - defer db.Close() - - if printConfig { - raw, err := install.RenderCopilotConfig(binaryPath, stateDir) - if err != nil { - return err - } - _, err = a.stdout.Write(append(raw, '\n')) - return err - } - - result, err := install.InstallCopilot(binaryPath, stateDir, copilotHome) - if err != nil { - return err - } - - fmt.Fprintf(a.stdout, "Installed Copilot hooks for Copilot CLI (sessionEnd) and VS Code (Stop).\nConfig: %s\nHook fragment: %s\nBinary: %s\nState dir: %s\nImport existing VS Code sessions with: arok capture --harness vscode --event scan\n", result.ConfigPath, result.FragmentPath, result.BinaryPath, result.StateDir) - return nil -} - +// runCapture dispatches to the harness-specific capture implementation. +// +// To add a new harness: +// 1. Add a case below calling a.runCapture(). +// 2. Create internal/cli/app_.go with that method. +// 3. Create internal// for the data parsing logic. +// +// See docs/adding-a-harness.md for a full walkthrough. func (a *App) runCapture(args []string) error { fs := flag.NewFlagSet("capture", flag.ContinueOnError) fs.SetOutput(a.stderr) @@ -166,261 +113,6 @@ func (a *App) runCapture(args []string) error { } } -func (a *App) runCaptureCopilot(eventName, stateDirOverride, payloadFile string, noReconcile bool) error { - if eventName == "" { - return errors.New("missing --event") - } - - stateDir, err := config.ResolveStateDir(stateDirOverride) - if err != nil { - return err - } - if err := config.EnsureLayout(stateDir); err != nil { - return err - } - - payloadRaw, err := readPayload(a.stdin, payloadFile) - if err != nil { - return err - } - payload, err := copilot.ParsePayload(payloadRaw) - if err != nil { - _ = appendLog(config.IngestLogPath(stateDir), fmt.Sprintf("%s capture failed: %v\n", time.Now().UTC().Format(time.RFC3339Nano), err)) - return err - } - if payload.SessionID() == "" { - return errors.New("Copilot payload is missing sessionId") - } - - if err := appendCaptureEvent(stateDir, eventName, payload); err != nil { - return err - } - - meta := metadataFromEnv() - sessionFile := copilot.ResolveSessionFile(payload) - summary, err := summarizeWithRetry(copilot.SummarizeOptions{ - EventName: eventName, - StateDir: stateDir, - Payload: payload, - SessionFile: sessionFile, - Meta: meta, - }, shutdownRetryAttempts(), shutdownRetryDelay()) - if err != nil { - _ = appendLog(config.IngestLogPath(stateDir), fmt.Sprintf("%s summarize failed for %s: %v\n", time.Now().UTC().Format(time.RFC3339Nano), payload.SessionID(), err)) - return err - } - - db, err := store.Open(stateDir) - if err != nil { - return err - } - defer db.Close() - - if err := db.UpsertSession(summary); err != nil { - _ = appendLog(config.IngestLogPath(stateDir), fmt.Sprintf("%s database write failed for %s: %v\n", time.Now().UTC().Format(time.RFC3339Nano), payload.SessionID(), err)) - return err - } - - if eventName == "sessionEnd" && summary.CaptureState != sessionpkg.CaptureStateFinal && !noReconcile { - if err := spawnDetachedReconcile(stateDir, payloadRaw, payload.SessionID(), eventName, sessionFile); err != nil { - _ = appendLog(config.ReconcileLogPath(stateDir), fmt.Sprintf("%s failed to schedule reconcile for %s: %v\n", time.Now().UTC().Format(time.RFC3339Nano), payload.SessionID(), err)) - return err - } - } - - return nil -} - -func (a *App) runCaptureVSCode(eventName, stateDirOverride, payloadFile string) error { - if eventName == "" { - return errors.New("missing --event") - } - - stateDir, err := config.ResolveStateDir(stateDirOverride) - if err != nil { - return err - } - if err := config.EnsureLayout(stateDir); err != nil { - return err - } - - db, err := store.Open(stateDir) - if err != nil { - return err - } - defer db.Close() - - if strings.EqualFold(eventName, "scan") { - return a.captureVSCodeScan(stateDir, db) - } - - payloadRaw, err := readPayload(a.stdin, payloadFile) - if err != nil { - return err - } - - payload, ok := parseVSCodeStopPayload(payloadRaw) - if !ok { - if len(strings.TrimSpace(string(payloadRaw))) > 0 { - _ = appendLog(config.IngestLogPath(stateDir), fmt.Sprintf("%s VS Code capture skipped: invalid Stop payload\n", time.Now().UTC().Format(time.RFC3339Nano))) - } - return nil - } - - sessionPath := deriveVSCodeChatSessionPath(payload.TranscriptPath) - if sessionPath == "" { - _ = appendLog(config.IngestLogPath(stateDir), fmt.Sprintf("%s VS Code capture skipped for %s: unable to derive chatSessions path\n", time.Now().UTC().Format(time.RFC3339Nano), payload.SessionID)) - return nil - } - - sessionData, err := vscode.ReadChatSession(sessionPath) - if err != nil { - _ = appendLog(config.IngestLogPath(stateDir), fmt.Sprintf("%s VS Code capture skipped for %s: %v\n", time.Now().UTC().Format(time.RFC3339Nano), payload.SessionID, err)) - return nil - } - if payload.SessionID != "" { - sessionData.SessionID = payload.SessionID - } - - summary := buildVSCodeSummary(sessionData, stateDir, eventName, sessionPath, payload.TranscriptPath) - if summary.SessionID == "" { - _ = appendLog(config.IngestLogPath(stateDir), fmt.Sprintf("%s VS Code capture skipped: missing session_id\n", time.Now().UTC().Format(time.RFC3339Nano))) - return nil - } - - if err := upsertVSCodeSummary(db, summary); err != nil { - _ = appendLog(config.IngestLogPath(stateDir), fmt.Sprintf("%s VS Code database write failed for %s: %v\n", time.Now().UTC().Format(time.RFC3339Nano), summary.SessionID, err)) - return nil - } - return nil -} - -func (a *App) captureVSCodeScan(stateDir string, db *store.Store) error { - userDataDir := vscode.DefaultUserDataDir() - if userDataDir == "" { - _ = appendLog(config.IngestLogPath(stateDir), fmt.Sprintf("%s VS Code scan skipped: user data directory not determinable\n", time.Now().UTC().Format(time.RFC3339Nano))) - return nil - } - - sessions, err := vscode.ScanSessions(userDataDir) - if err != nil { - _ = appendLog(config.IngestLogPath(stateDir), fmt.Sprintf("%s VS Code scan skipped: %v\n", time.Now().UTC().Format(time.RFC3339Nano), err)) - return nil - } - - for _, sessionData := range sessions { - if !shouldImportScannedVSCodeSession(sessionData) { - continue - } - summary := buildVSCodeSummary(sessionData, stateDir, "scan", "", "") - if summary.SessionID == "" { - continue - } - if err := upsertVSCodeSummary(db, summary); err != nil { - _ = appendLog(config.IngestLogPath(stateDir), fmt.Sprintf("%s VS Code scan write failed for %s: %v\n", time.Now().UTC().Format(time.RFC3339Nano), summary.SessionID, err)) - } - } - return nil -} - -func (a *App) runReconcile(args []string) error { - fs := flag.NewFlagSet("reconcile", flag.ContinueOnError) - fs.SetOutput(a.stderr) - var ( - harness string - eventName string - stateDirOverride string - payloadFile string - sessionFile string - sessionID string - attempts int - delay time.Duration - initialDelay time.Duration - ) - fs.StringVar(&harness, "harness", "", "Harness name.") - fs.StringVar(&eventName, "event", "sessionEnd", "Hook event name.") - fs.StringVar(&stateDirOverride, "state-dir", "", "Override the AROK state directory.") - fs.StringVar(&payloadFile, "payload-file", "", "Persisted hook payload file.") - fs.StringVar(&sessionFile, "session-file", "", "Explicit session log path.") - fs.StringVar(&sessionID, "session-id", "", "Session identifier.") - fs.IntVar(&attempts, "attempts", asyncReconcileAttempts(), "Number of reconciliation attempts.") - fs.DurationVar(&delay, "delay", asyncReconcileDelay(), "Delay between reconciliation attempts.") - fs.DurationVar(&initialDelay, "initial-delay", asyncReconcileInitialDelay(), "Delay before the first reconciliation attempt.") - if err := fs.Parse(args); err != nil { - return err - } - - if harness != "copilot" { - return fmt.Errorf("unsupported harness %q", harness) - } - if payloadFile == "" || sessionFile == "" || sessionID == "" { - return errors.New("reconcile requires --payload-file, --session-file, and --session-id") - } - - stateDir, err := config.ResolveStateDir(stateDirOverride) - if err != nil { - return err - } - if err := config.EnsureLayout(stateDir); err != nil { - return err - } - - payloadRaw, err := os.ReadFile(payloadFile) - if err != nil { - return fmt.Errorf("read payload snapshot: %w", err) - } - payload, err := copilot.ParsePayload(payloadRaw) - if err != nil { - return err - } - - if initialDelay > 0 { - time.Sleep(initialDelay) - } - - meta := metadataFromEnv() - summary, err := summarizeWithRetry(copilot.SummarizeOptions{ - EventName: eventName, - StateDir: stateDir, - Payload: payload, - SessionFile: sessionFile, - Meta: meta, - }, attempts, delay) - if err != nil { - return err - } - - db, err := store.Open(stateDir) - if err != nil { - return err - } - defer db.Close() - - // If reconcile exhausted without finding shutdown metrics, mark as best_effort - // rather than leaving it provisional forever. This is expected for sessions that - // end without emitting session.shutdown.modelMetrics (e.g. abrupt exits). - reconcileExhausted := summary.CaptureState != sessionpkg.CaptureStateFinal - if reconcileExhausted { - summary.CaptureState = sessionpkg.CaptureStateBestEffort - } - - if err := db.UpsertSession(summary); err != nil { - return err - } - - if err := os.Remove(payloadFile); err != nil && !os.IsNotExist(err) { - _ = appendLog(config.ReconcileLogPath(stateDir), fmt.Sprintf("%s cleanup failed for %s: %v\n", time.Now().UTC().Format(time.RFC3339Nano), sessionID, err)) - } - - if reconcileExhausted { - message := fmt.Sprintf("%s reconcile exhausted before final totals for %s; marked best_effort\n", time.Now().UTC().Format(time.RFC3339Nano), sessionID) - _ = appendLog(config.ReconcileLogPath(stateDir), message) - return errors.New(strings.TrimSpace(message)) - } - - return nil -} - func (a *App) runQuery(args []string) error { subcommand := "sessions" if len(args) > 0 { @@ -884,46 +576,11 @@ func (a *App) printRootUsage() { fmt.Fprintf(a.stdout, "arok %s\n\nCommands:\n install copilot\n capture --harness [copilot|vscode] --event \n reconcile --harness copilot\n query [sessions|hosts|repos|branches|worktrees|harnesses|tasks|models]\n analyze [overview|missing-finals]\n doctor\n update\n version\n", version.Version) } -func summarizeWithRetry(opts copilot.SummarizeOptions, attempts int, delay time.Duration) (sessionpkg.SessionSummary, error) { - if attempts < 1 { - attempts = 1 - } - - var summary sessionpkg.SessionSummary - var err error - for attempt := 1; attempt <= attempts; attempt++ { - summary, err = copilot.Summarize(opts) - if err != nil { - return sessionpkg.SessionSummary{}, err - } - if opts.EventName != "sessionEnd" || summary.CaptureState == sessionpkg.CaptureStateFinal || attempt == attempts { - return summary, nil - } - time.Sleep(delay) - } - return summary, nil -} - -func metadataFromEnv() copilot.Meta { - return copilot.Meta{ - TaskID: strings.TrimSpace(os.Getenv("LLM_USAGE_TASK_ID")), - Title: strings.TrimSpace(os.Getenv("LLM_USAGE_TITLE")), - Summary: strings.TrimSpace(os.Getenv("LLM_USAGE_SUMMARY")), - Tags: splitCSV(os.Getenv("LLM_USAGE_TAGS")), - } -} - -func appendCaptureEvent(stateDir, eventName string, payload copilot.Payload) error { - record := map[string]any{ - "captured_at": time.Now().UTC().Format(time.RFC3339Nano), - "event_name": eventName, - "payload": payload, - } - raw, err := json.Marshal(record) - if err != nil { - return err +func readPayload(stdin io.Reader, payloadFile string) ([]byte, error) { + if payloadFile != "" { + return os.ReadFile(payloadFile) } - return appendLog(config.CaptureLogPath(stateDir), string(raw)+"\n") + return io.ReadAll(stdin) } func appendLog(path, line string) error { @@ -936,57 +593,6 @@ func appendLog(path, line string) error { return err } -func spawnDetachedReconcile(stateDir string, payloadRaw []byte, sessionID, eventName, sessionFile string) error { - snapshot, err := os.CreateTemp(config.ReconcileDir(stateDir), sessionID+".payload.*.json") - if err != nil { - return err - } - if _, err := snapshot.Write(payloadRaw); err != nil { - snapshot.Close() - return err - } - if err := snapshot.Close(); err != nil { - return err - } - - exe, err := os.Executable() - if err != nil { - return err - } - - logFile, err := os.OpenFile(config.ReconcileLogPath(stateDir), os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0o644) - if err != nil { - return err - } - defer logFile.Close() - - cmd := exec.Command( - exe, - "reconcile", - "--harness", "copilot", - "--event", eventName, - "--state-dir", stateDir, - "--payload-file", snapshot.Name(), - "--session-file", sessionFile, - "--session-id", sessionID, - "--attempts", fmt.Sprintf("%d", asyncReconcileAttempts()), - "--delay", asyncReconcileDelay().String(), - "--initial-delay", asyncReconcileInitialDelay().String(), - ) - cmd.Stdin = nil - cmd.Stdout = logFile - cmd.Stderr = logFile - cmd.SysProcAttr = &syscall.SysProcAttr{Setsid: true} - return cmd.Start() -} - -func readPayload(stdin io.Reader, payloadFile string) ([]byte, error) { - if payloadFile != "" { - return os.ReadFile(payloadFile) - } - return io.ReadAll(stdin) -} - func parseSince(raw string) (*time.Time, error) { if raw == "" { return nil, nil @@ -1086,64 +692,6 @@ func writeSessionTable(w io.Writer, title string, rows []sessionpkg.SessionListI fmt.Fprintln(w) } -func splitCSV(raw string) []string { - if strings.TrimSpace(raw) == "" { - return nil - } - parts := strings.Split(raw, ",") - out := make([]string, 0, len(parts)) - for _, part := range parts { - if trimmed := strings.TrimSpace(part); trimmed != "" { - out = append(out, trimmed) - } - } - return out -} - -func shutdownRetryAttempts() int { - return envInt("AROK_COPILOT_SHUTDOWN_RETRY_ATTEMPTS", 6) -} - -func shutdownRetryDelay() time.Duration { - return envDuration("AROK_COPILOT_SHUTDOWN_RETRY_DELAY", 500*time.Millisecond) -} - -func asyncReconcileAttempts() int { - return envInt("AROK_COPILOT_RECONCILE_ATTEMPTS", 12) -} - -func asyncReconcileDelay() time.Duration { - return envDuration("AROK_COPILOT_RECONCILE_DELAY", time.Second) -} - -func asyncReconcileInitialDelay() time.Duration { - return envDuration("AROK_COPILOT_RECONCILE_INITIAL_DELAY", 2*time.Second) -} - -func envInt(key string, fallback int) int { - value := strings.TrimSpace(os.Getenv(key)) - if value == "" { - return fallback - } - var parsed int - if _, err := fmt.Sscanf(value, "%d", &parsed); err != nil || parsed < 1 { - return fallback - } - return parsed -} - -func envDuration(key string, fallback time.Duration) time.Duration { - value := strings.TrimSpace(os.Getenv(key)) - if value == "" { - return fallback - } - parsed, err := time.ParseDuration(value) - if err != nil || parsed < 0 { - return fallback - } - return parsed -} - func statusString(ok bool) string { if ok { return "ok" @@ -1151,168 +699,6 @@ func statusString(ok bool) string { return "missing" } -type vscodeStopPayload struct { - HookEventName string `json:"hook_event_name"` - SessionID string `json:"session_id"` - TranscriptPath string `json:"transcript_path"` - CWD string `json:"cwd"` -} - -func parseVSCodeStopPayload(raw []byte) (vscodeStopPayload, bool) { - if len(strings.TrimSpace(string(raw))) == 0 { - return vscodeStopPayload{}, false - } - - var payload vscodeStopPayload - if err := json.Unmarshal(raw, &payload); err != nil { - return vscodeStopPayload{}, false - } - if payload.SessionID == "" || payload.TranscriptPath == "" { - return vscodeStopPayload{}, false - } - return payload, true -} - -func deriveVSCodeChatSessionPath(transcriptPath string) string { - if strings.TrimSpace(transcriptPath) == "" { - return "" - } - sessionFile := filepath.Base(transcriptPath) - storageDir := filepath.Dir(filepath.Dir(filepath.Dir(transcriptPath))) - if sessionFile == "." || sessionFile == string(filepath.Separator) { - return "" - } - return filepath.Join(storageDir, "chatSessions", sessionFile) -} - -func buildVSCodeSummary(sessionData vscode.Session, stateDir, eventName, eventLogPath, transcriptPath string) sessionpkg.SessionSummary { - hostName, _ := os.Hostname() - git := gitmeta.Inspect(sessionData.WorkspaceFolder) - - totalOutputTokens := int64(0) - totalInputTokens := int64(0) - modelStats := map[string]int64{} - modelCounts := map[string]int64{} - endedAt := sessionData.CreationDate - for _, req := range sessionData.Requests { - totalOutputTokens += req.CompletionTokens - totalInputTokens += req.PromptTokens - model := req.ModelID - if model == "" { - model = "unknown" - } - modelStats[model] += req.CompletionTokens - modelCounts[model]++ - if req.Timestamp.After(endedAt) { - endedAt = req.Timestamp - } - } - - modelNames := make([]string, 0, len(modelStats)) - for model := range modelStats { - modelNames = append(modelNames, model) - } - slices.Sort(modelNames) - - models := make([]sessionpkg.ModelUsage, 0, len(modelNames)) - for _, model := range modelNames { - tokens := modelStats[model] - count := modelCounts[model] - models = append(models, sessionpkg.ModelUsage{ - Model: model, - AssistantMessageCount: count, - AssistantOutputTokens: tokens, - OutputTokens: sessionpkg.PtrInt64(tokens), - RequestCount: sessionpkg.PtrInt64(count), - }) - } - - startedAt := "" - if !sessionData.CreationDate.IsZero() { - startedAt = sessionData.CreationDate.UTC().Format(time.RFC3339) - } - endedAtRaw := "" - if !endedAt.IsZero() { - endedAtRaw = endedAt.UTC().Format(time.RFC3339) - } - - notes := []string{ - "VS Code Copilot usage is summarized from the local chatSessions JSONL transaction log.", - "Token counts are sourced from result.metadata (primary) or completionTokens/usage fields (fallback).", - } - if sessionData.WorkspaceFolder == "" { - notes = append(notes, "Workspace metadata is unavailable for missing or remote VS Code workspaces.") - } - - var totalInputTokensPtr *int64 - if totalInputTokens > 0 { - totalInputTokensPtr = sessionpkg.PtrInt64(totalInputTokens) - } - - return sessionpkg.SessionSummary{ - SchemaVersion: 1, - Source: "vscode", - Harness: sessionpkg.HarnessVSCodeCopilot, - CollectedAt: time.Now().UTC().Format(time.RFC3339Nano), - SessionID: sessionData.SessionID, - EventName: eventName, - CaptureState: sessionpkg.CaptureStateFinal, - UsageSource: "chatSessions.completionTokens", - TranscriptPath: transcriptPath, - EventLogPath: eventLogPath, - StateDir: stateDir, - CWD: sessionData.WorkspaceFolder, - RepoRoot: git.RepoRoot, - WorktreeRoot: git.WorktreeRoot, - GitCommonDir: git.GitCommonDir, - RepoRemote: git.RepoRemote, - RepoBranch: git.RepoBranch, - RepoHead: git.RepoHead, - HostName: hostName, - StartedAt: startedAt, - EndedAt: endedAtRaw, - InteractionCount: int64(len(sessionData.Requests)), - AssistantMessageCount: int64(len(sessionData.Requests)), - AssistantOutputTokens: totalOutputTokens, - TotalInputTokens: totalInputTokensPtr, - TotalOutputTokens: sessionpkg.PtrInt64(totalOutputTokens), - Models: models, - Notes: notes, - } -} - -func shouldImportScannedVSCodeSession(sessionData vscode.Session) bool { - if len(sessionData.Requests) == 0 { - return true - } - for _, req := range sessionData.Requests { - if req.CompletionTokens > 0 { - return true - } - } - return false -} - -func upsertVSCodeSummary(db *store.Store, summary sessionpkg.SessionSummary) error { - existing, err := db.GetSession(summary.SessionID) - if err != nil && !errors.Is(err, store.ErrSessionNotFound) { - return err - } - if err == nil && existing.CaptureState == sessionpkg.CaptureStateFinal && - derefInt64(existing.TotalOutputTokens) == derefInt64(summary.TotalOutputTokens) && - derefInt64(existing.TotalInputTokens) == derefInt64(summary.TotalInputTokens) { - return nil - } - return db.UpsertSession(summary) -} - -func derefInt64(value *int64) int64 { - if value == nil { - return 0 - } - return *value -} - func displayString(value, fallback string) string { if strings.TrimSpace(value) == "" { return fallback diff --git a/internal/cli/app_copilot.go b/internal/cli/app_copilot.go new file mode 100644 index 0000000..d69eb04 --- /dev/null +++ b/internal/cli/app_copilot.go @@ -0,0 +1,396 @@ +package cli + +// Copilot CLI harness — capture, reconcile, and install support. +// +// To add a new harness, follow the same pattern as this file: +// 1. Create internal/cli/app_.go with a runCapture() method and helpers. +// 2. Create internal// for the payload parsing and summarizing logic. +// 3. Add a case in runCapture() in app.go that calls runCapture(). +// 4. Optionally add runInstall() and a case in runInstall(). +// +// See docs/adding-a-harness.md for the full walkthrough. + +import ( + "encoding/json" + "errors" + "flag" + "fmt" + "os" + "os/exec" + "strings" + "syscall" + "time" + + "github.com/srbouffard/arok/internal/config" + "github.com/srbouffard/arok/internal/copilot" + "github.com/srbouffard/arok/internal/install" + sessionpkg "github.com/srbouffard/arok/internal/session" + "github.com/srbouffard/arok/internal/store" +) + +// runInstallCopilot handles: arok install copilot [flags] +func (a *App) runInstallCopilot(args []string) error { + fs := flag.NewFlagSet("install copilot", flag.ContinueOnError) + fs.SetOutput(a.stderr) + var ( + stateDirOverride string + copilotHome = install.DefaultCopilotHome() + binaryPath string + printConfig bool + ) + fs.StringVar(&stateDirOverride, "state-dir", "", "Override the AROK state directory.") + fs.StringVar(&copilotHome, "copilot-home", copilotHome, "Override the Copilot home directory.") + fs.StringVar(&binaryPath, "binary-path", "", "Override the binary path written into the Copilot hook config.") + fs.BoolVar(&printConfig, "print-config", false, "Print the generated Copilot hook config instead of installing it.") + if err := fs.Parse(args); err != nil { + return err + } + + stateDir, err := config.ResolveStateDir(stateDirOverride) + if err != nil { + return err + } + if err := config.EnsureLayout(stateDir); err != nil { + return err + } + + if binaryPath == "" { + binaryPath, err = os.Executable() + if err != nil { + return fmt.Errorf("resolve executable path: %w", err) + } + } + + db, err := store.Open(stateDir) + if err != nil { + return err + } + defer db.Close() + + if printConfig { + raw, err := install.RenderCopilotConfig(binaryPath, stateDir) + if err != nil { + return err + } + _, err = a.stdout.Write(append(raw, '\n')) + return err + } + + result, err := install.InstallCopilot(binaryPath, stateDir, copilotHome) + if err != nil { + return err + } + + fmt.Fprintf(a.stdout, "Installed Copilot hooks for Copilot CLI (sessionEnd) and VS Code (Stop).\nConfig: %s\nHook fragment: %s\nBinary: %s\nState dir: %s\nImport existing VS Code sessions with: arok capture --harness vscode --event scan\n", result.ConfigPath, result.FragmentPath, result.BinaryPath, result.StateDir) + return nil +} + +// runCaptureCopilot handles: arok capture --harness copilot --event +func (a *App) runCaptureCopilot(eventName, stateDirOverride, payloadFile string, noReconcile bool) error { + if eventName == "" { + return errors.New("missing --event") + } + + stateDir, err := config.ResolveStateDir(stateDirOverride) + if err != nil { + return err + } + if err := config.EnsureLayout(stateDir); err != nil { + return err + } + + payloadRaw, err := readPayload(a.stdin, payloadFile) + if err != nil { + return err + } + payload, err := copilot.ParsePayload(payloadRaw) + if err != nil { + _ = appendLog(config.IngestLogPath(stateDir), fmt.Sprintf("%s capture failed: %v\n", time.Now().UTC().Format(time.RFC3339Nano), err)) + return err + } + if payload.SessionID() == "" { + return errors.New("Copilot payload is missing sessionId") + } + + if err := appendCaptureEvent(stateDir, eventName, payload); err != nil { + return err + } + + meta := metadataFromEnv() + sessionFile := copilot.ResolveSessionFile(payload) + summary, err := summarizeWithRetry(copilot.SummarizeOptions{ + EventName: eventName, + StateDir: stateDir, + Payload: payload, + SessionFile: sessionFile, + Meta: meta, + }, shutdownRetryAttempts(), shutdownRetryDelay()) + if err != nil { + _ = appendLog(config.IngestLogPath(stateDir), fmt.Sprintf("%s summarize failed for %s: %v\n", time.Now().UTC().Format(time.RFC3339Nano), payload.SessionID(), err)) + return err + } + + db, err := store.Open(stateDir) + if err != nil { + return err + } + defer db.Close() + + if err := db.UpsertSession(summary); err != nil { + _ = appendLog(config.IngestLogPath(stateDir), fmt.Sprintf("%s database write failed for %s: %v\n", time.Now().UTC().Format(time.RFC3339Nano), payload.SessionID(), err)) + return err + } + + if eventName == "sessionEnd" && summary.CaptureState != sessionpkg.CaptureStateFinal && !noReconcile { + if err := spawnDetachedReconcile(stateDir, payloadRaw, payload.SessionID(), eventName, sessionFile); err != nil { + _ = appendLog(config.ReconcileLogPath(stateDir), fmt.Sprintf("%s failed to schedule reconcile for %s: %v\n", time.Now().UTC().Format(time.RFC3339Nano), payload.SessionID(), err)) + return err + } + } + + return nil +} + +// runReconcile handles: arok reconcile --harness copilot [flags] +// +// Reconcile is a copilot-specific background process: it polls events.jsonl until +// session.shutdown emits final modelMetrics, then upgrades the stored record from +// provisional to final. Harnesses that produce final data at capture time do not +// need this command. +func (a *App) runReconcile(args []string) error { + fs := flag.NewFlagSet("reconcile", flag.ContinueOnError) + fs.SetOutput(a.stderr) + var ( + harness string + eventName string + stateDirOverride string + payloadFile string + sessionFile string + sessionID string + attempts int + delay time.Duration + initialDelay time.Duration + ) + fs.StringVar(&harness, "harness", "", "Harness name.") + fs.StringVar(&eventName, "event", "sessionEnd", "Hook event name.") + fs.StringVar(&stateDirOverride, "state-dir", "", "Override the AROK state directory.") + fs.StringVar(&payloadFile, "payload-file", "", "Persisted hook payload file.") + fs.StringVar(&sessionFile, "session-file", "", "Explicit session log path.") + fs.StringVar(&sessionID, "session-id", "", "Session identifier.") + fs.IntVar(&attempts, "attempts", asyncReconcileAttempts(), "Number of reconciliation attempts.") + fs.DurationVar(&delay, "delay", asyncReconcileDelay(), "Delay between reconciliation attempts.") + fs.DurationVar(&initialDelay, "initial-delay", asyncReconcileInitialDelay(), "Delay before the first reconciliation attempt.") + if err := fs.Parse(args); err != nil { + return err + } + + if harness != "copilot" { + return fmt.Errorf("unsupported harness %q", harness) + } + if payloadFile == "" || sessionFile == "" || sessionID == "" { + return errors.New("reconcile requires --payload-file, --session-file, and --session-id") + } + + stateDir, err := config.ResolveStateDir(stateDirOverride) + if err != nil { + return err + } + if err := config.EnsureLayout(stateDir); err != nil { + return err + } + + payloadRaw, err := os.ReadFile(payloadFile) + if err != nil { + return fmt.Errorf("read payload snapshot: %w", err) + } + payload, err := copilot.ParsePayload(payloadRaw) + if err != nil { + return err + } + + if initialDelay > 0 { + time.Sleep(initialDelay) + } + + meta := metadataFromEnv() + summary, err := summarizeWithRetry(copilot.SummarizeOptions{ + EventName: eventName, + StateDir: stateDir, + Payload: payload, + SessionFile: sessionFile, + Meta: meta, + }, attempts, delay) + if err != nil { + return err + } + + db, err := store.Open(stateDir) + if err != nil { + return err + } + defer db.Close() + + // If reconcile exhausted without finding shutdown metrics, mark as best_effort + // rather than leaving it provisional forever. This is expected for sessions that + // end without emitting session.shutdown.modelMetrics (e.g. abrupt exits). + reconcileExhausted := summary.CaptureState != sessionpkg.CaptureStateFinal + if reconcileExhausted { + summary.CaptureState = sessionpkg.CaptureStateBestEffort + } + + if err := db.UpsertSession(summary); err != nil { + return err + } + + if err := os.Remove(payloadFile); err != nil && !os.IsNotExist(err) { + _ = appendLog(config.ReconcileLogPath(stateDir), fmt.Sprintf("%s cleanup failed for %s: %v\n", time.Now().UTC().Format(time.RFC3339Nano), sessionID, err)) + } + + if reconcileExhausted { + message := fmt.Sprintf("%s reconcile exhausted before final totals for %s; marked best_effort\n", time.Now().UTC().Format(time.RFC3339Nano), sessionID) + _ = appendLog(config.ReconcileLogPath(stateDir), message) + return errors.New(strings.TrimSpace(message)) + } + + return nil +} + +func summarizeWithRetry(opts copilot.SummarizeOptions, attempts int, delay time.Duration) (sessionpkg.SessionSummary, error) { + if attempts < 1 { + attempts = 1 + } + var ( + summary sessionpkg.SessionSummary + err error + ) + for attempt := 1; attempt <= attempts; attempt++ { + summary, err = copilot.Summarize(opts) + if err != nil { + return sessionpkg.SessionSummary{}, err + } + if opts.EventName != "sessionEnd" || summary.CaptureState == sessionpkg.CaptureStateFinal || attempt == attempts { + return summary, nil + } + time.Sleep(delay) + } + return summary, nil +} + +func metadataFromEnv() copilot.Meta { + return copilot.Meta{ + TaskID: strings.TrimSpace(os.Getenv("LLM_USAGE_TASK_ID")), + Title: strings.TrimSpace(os.Getenv("LLM_USAGE_TITLE")), + Summary: strings.TrimSpace(os.Getenv("LLM_USAGE_SUMMARY")), + Tags: splitCSV(os.Getenv("LLM_USAGE_TAGS")), + } +} + +func appendCaptureEvent(stateDir, eventName string, payload copilot.Payload) error { + record := map[string]any{ + "captured_at": time.Now().UTC().Format(time.RFC3339Nano), + "event_name": eventName, + "payload": payload, + } + raw, err := json.Marshal(record) + if err != nil { + return err + } + return appendLog(config.CaptureLogPath(stateDir), string(raw)+"\n") +} + +func spawnDetachedReconcile(stateDir string, payloadRaw []byte, sessionID, eventName, sessionFile string) error { + snapshot, err := os.CreateTemp(config.ReconcileDir(stateDir), sessionID+".payload.*.json") + if err != nil { + return err + } + if _, err := snapshot.Write(payloadRaw); err != nil { + snapshot.Close() + return err + } + if err := snapshot.Close(); err != nil { + return err + } + + exe, err := os.Executable() + if err != nil { + return err + } + + logFile, err := os.OpenFile(config.ReconcileLogPath(stateDir), os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0o644) + if err != nil { + return err + } + // Close the parent's copy of the fd after cmd.Start() hands it to the child. + // The child process writes directly to the kernel fd, so no Go-side buffering + // exists to lose. The _ discard is intentional. + defer func() { _ = logFile.Close() }() + + cmd := exec.Command( + exe, + "reconcile", + "--harness", "copilot", + "--event", eventName, + "--state-dir", stateDir, + "--payload-file", snapshot.Name(), + "--session-file", sessionFile, + "--session-id", sessionID, + "--attempts", fmt.Sprintf("%d", asyncReconcileAttempts()), + "--delay", asyncReconcileDelay().String(), + "--initial-delay", asyncReconcileInitialDelay().String(), + ) + cmd.Stdin = nil + cmd.Stdout = logFile + cmd.Stderr = logFile + cmd.SysProcAttr = &syscall.SysProcAttr{Setsid: true} + return cmd.Start() +} + +func shutdownRetryAttempts() int { return envInt("AROK_COPILOT_SHUTDOWN_RETRY_ATTEMPTS", 6) } +func shutdownRetryDelay() time.Duration { + return envDuration("AROK_COPILOT_SHUTDOWN_RETRY_DELAY", 500*time.Millisecond) +} +func asyncReconcileAttempts() int { return envInt("AROK_COPILOT_RECONCILE_ATTEMPTS", 12) } +func asyncReconcileDelay() time.Duration { + return envDuration("AROK_COPILOT_RECONCILE_DELAY", time.Second) +} +func asyncReconcileInitialDelay() time.Duration { + return envDuration("AROK_COPILOT_RECONCILE_INITIAL_DELAY", 2*time.Second) +} + +func envInt(key string, fallback int) int { + value := strings.TrimSpace(os.Getenv(key)) + if value == "" { + return fallback + } + var parsed int + if _, err := fmt.Sscanf(value, "%d", &parsed); err != nil || parsed < 1 { + return fallback + } + return parsed +} + +func envDuration(key string, fallback time.Duration) time.Duration { + value := strings.TrimSpace(os.Getenv(key)) + if value == "" { + return fallback + } + parsed, err := time.ParseDuration(value) + if err != nil || parsed < 0 { + return fallback + } + return parsed +} + +func splitCSV(raw string) []string { + if strings.TrimSpace(raw) == "" { + return nil + } + parts := strings.Split(raw, ",") + out := make([]string, 0, len(parts)) + for _, part := range parts { + if trimmed := strings.TrimSpace(part); trimmed != "" { + out = append(out, trimmed) + } + } + return out +} diff --git a/internal/cli/app_copilot_test.go b/internal/cli/app_copilot_test.go new file mode 100644 index 0000000..e1909e7 --- /dev/null +++ b/internal/cli/app_copilot_test.go @@ -0,0 +1,322 @@ +package cli + +import ( + "bytes" + "encoding/json" + "fmt" + "os" + "path/filepath" + "strings" + "testing" + + sessionpkg "github.com/srbouffard/arok/internal/session" + "github.com/srbouffard/arok/internal/store" +) + +// sessionFinalJSONL is a minimal events.jsonl containing a session.shutdown event +// with modelMetrics — the authoritative source for final token totals. +const sessionFinalJSONL = "" + + `{"timestamp":"2026-01-01T00:00:00Z","type":"assistant.message","data":{"model":"claude-sonnet","outputTokens":150,"interactionId":"i-1","toolRequests":[{}]}}` + "\n" + + `{"timestamp":"2026-01-01T00:00:01Z","type":"tool.execution_complete","data":{"success":true}}` + "\n" + + `{"timestamp":"2026-01-01T00:00:02Z","type":"session.shutdown","data":{"modelMetrics":{"claude-sonnet":{"usage":{"inputTokens":500,"outputTokens":150,"cacheReadTokens":200},"requests":{"count":1}}}}}` + "\n" + +// sessionProvisionalJSONL is a minimal events.jsonl without a session.shutdown event, +// simulating a session that ended before metrics were written. +const sessionProvisionalJSONL = "" + + `{"timestamp":"2026-01-01T00:00:00Z","type":"assistant.message","data":{"model":"claude-sonnet","outputTokens":150,"interactionId":"i-1","toolRequests":[{}]}}` + "\n" + + `{"timestamp":"2026-01-01T00:00:01Z","type":"tool.execution_complete","data":{"success":true}}` + "\n" + +func TestRunCaptureCopilotFinalSession(t *testing.T) { + stateDir := t.TempDir() + copilotHome := t.TempDir() + sessionID := "sess-copilot-final" + + eventsDir := filepath.Join(copilotHome, "session-state", sessionID) + if err := os.MkdirAll(eventsDir, 0o755); err != nil { + t.Fatalf("MkdirAll(eventsDir) returned error: %v", err) + } + eventsFile := filepath.Join(eventsDir, "events.jsonl") + if err := os.WriteFile(eventsFile, []byte(sessionFinalJSONL), 0o644); err != nil { + t.Fatalf("WriteFile(events.jsonl) returned error: %v", err) + } + + payload := fmt.Sprintf(`{"sessionId":%q,"cwd":%q}`, sessionID, t.TempDir()) + payloadFile := writeTempFile(t, "payload.json", payload) + t.Setenv("COPILOT_HOME", copilotHome) + + app := New(bytes.NewReader(nil), &bytes.Buffer{}, &bytes.Buffer{}) + if err := app.Run([]string{ + "capture", "--harness", "copilot", "--event", "sessionEnd", + "--state-dir", stateDir, + "--payload-file", payloadFile, + "--no-reconcile", + }); err != nil { + t.Fatalf("Run returned error: %v", err) + } + + db, err := store.Open(stateDir) + if err != nil { + t.Fatalf("Open returned error: %v", err) + } + defer db.Close() + + summary, err := db.GetSession(sessionID) + if err != nil { + t.Fatalf("GetSession returned error: %v", err) + } + if summary.Harness != sessionpkg.HarnessCopilotCLI { + t.Errorf("Harness = %q, want %q", summary.Harness, sessionpkg.HarnessCopilotCLI) + } + if summary.CaptureState != sessionpkg.CaptureStateFinal { + t.Errorf("CaptureState = %q, want %q", summary.CaptureState, sessionpkg.CaptureStateFinal) + } + if summary.TotalInputTokens == nil || *summary.TotalInputTokens != 500 { + t.Errorf("TotalInputTokens = %#v, want 500", summary.TotalInputTokens) + } + if summary.TotalOutputTokens == nil || *summary.TotalOutputTokens != 150 { + t.Errorf("TotalOutputTokens = %#v, want 150", summary.TotalOutputTokens) + } + if summary.TotalCacheReadTokens == nil || *summary.TotalCacheReadTokens != 200 { + t.Errorf("TotalCacheReadTokens = %#v, want 200", summary.TotalCacheReadTokens) + } + if summary.EventName != "sessionEnd" { + t.Errorf("EventName = %q, want sessionEnd", summary.EventName) + } +} + +func TestRunCaptureCopilotProvisionalSession(t *testing.T) { + stateDir := t.TempDir() + copilotHome := t.TempDir() + sessionID := "sess-copilot-provisional" + + eventsDir := filepath.Join(copilotHome, "session-state", sessionID) + if err := os.MkdirAll(eventsDir, 0o755); err != nil { + t.Fatalf("MkdirAll(eventsDir) returned error: %v", err) + } + eventsFile := filepath.Join(eventsDir, "events.jsonl") + if err := os.WriteFile(eventsFile, []byte(sessionProvisionalJSONL), 0o644); err != nil { + t.Fatalf("WriteFile(events.jsonl) returned error: %v", err) + } + + payload := fmt.Sprintf(`{"sessionId":%q,"cwd":%q}`, sessionID, t.TempDir()) + payloadFile := writeTempFile(t, "payload.json", payload) + t.Setenv("COPILOT_HOME", copilotHome) + // Disable retry so the test doesn't spin waiting for shutdown metrics. + t.Setenv("AROK_COPILOT_SHUTDOWN_RETRY_ATTEMPTS", "1") + + app := New(bytes.NewReader(nil), &bytes.Buffer{}, &bytes.Buffer{}) + if err := app.Run([]string{ + "capture", "--harness", "copilot", "--event", "sessionEnd", + "--state-dir", stateDir, + "--payload-file", payloadFile, + "--no-reconcile", + }); err != nil { + t.Fatalf("Run returned error: %v", err) + } + + db, err := store.Open(stateDir) + if err != nil { + t.Fatalf("Open returned error: %v", err) + } + defer db.Close() + + summary, err := db.GetSession(sessionID) + if err != nil { + t.Fatalf("GetSession returned error: %v", err) + } + if summary.CaptureState != sessionpkg.CaptureStateProvisional { + t.Errorf("CaptureState = %q, want %q", summary.CaptureState, sessionpkg.CaptureStateProvisional) + } + // Without shutdown metrics, TotalOutputTokens falls back to assistant.message counts. + if summary.TotalOutputTokens == nil || *summary.TotalOutputTokens != 150 { + t.Errorf("TotalOutputTokens = %#v, want 150 (fallback from assistant.message)", summary.TotalOutputTokens) + } +} + +func TestRunCaptureCopilotMissingSessionID(t *testing.T) { + stateDir := t.TempDir() + payloadFile := writeTempFile(t, "payload.json", `{"cwd":"/tmp"}`) + + app := New(bytes.NewReader(nil), &bytes.Buffer{}, &bytes.Buffer{}) + err := app.Run([]string{ + "capture", "--harness", "copilot", "--event", "sessionEnd", + "--state-dir", stateDir, + "--payload-file", payloadFile, + "--no-reconcile", + }) + if err == nil { + t.Fatal("expected error for missing sessionId, got nil") + } +} + +func TestRunCaptureCopilotUnknownHarness(t *testing.T) { + app := New(bytes.NewReader(nil), &bytes.Buffer{}, &bytes.Buffer{}) + err := app.Run([]string{"capture", "--harness", "unknown-harness", "--event", "sessionEnd"}) + if err == nil { + t.Fatal("expected error for unknown harness, got nil") + } +} + +func TestRunInstallCopilotPrintConfig(t *testing.T) { + stateDir := t.TempDir() + var out bytes.Buffer + app := New(bytes.NewReader(nil), &out, &bytes.Buffer{}) + if err := app.Run([]string{ + "install", "copilot", + "--print-config", + "--state-dir", stateDir, + "--binary-path", "/usr/local/bin/arok", + }); err != nil { + t.Fatalf("Run returned error: %v", err) + } + output := out.String() + if !strings.Contains(output, "sessionEnd") { + t.Error("expected config to contain sessionEnd hook") + } + if !strings.Contains(output, "/usr/local/bin/arok") { + t.Error("expected config to contain binary path") + } + var parsed map[string]any + if err := json.Unmarshal([]byte(output), &parsed); err != nil { + t.Errorf("config is not valid JSON: %v", err) + } +} + +func TestRunReconcileUpgradesProvisionalToFinal(t *testing.T) { + stateDir := t.TempDir() + copilotHome := t.TempDir() + sessionID := "sess-reconcile-final" + + eventsDir := filepath.Join(copilotHome, "session-state", sessionID) + if err := os.MkdirAll(eventsDir, 0o755); err != nil { + t.Fatalf("MkdirAll returned error: %v", err) + } + eventsFile := filepath.Join(eventsDir, "events.jsonl") + + // Stage 1: capture with no shutdown event → provisional. + if err := os.WriteFile(eventsFile, []byte(sessionProvisionalJSONL), 0o644); err != nil { + t.Fatalf("WriteFile(provisional) returned error: %v", err) + } + payload := fmt.Sprintf(`{"sessionId":%q,"cwd":%q}`, sessionID, t.TempDir()) + payloadFile := writeTempFile(t, "payload.json", payload) + t.Setenv("COPILOT_HOME", copilotHome) + t.Setenv("AROK_COPILOT_SHUTDOWN_RETRY_ATTEMPTS", "1") + + app := New(bytes.NewReader(nil), &bytes.Buffer{}, &bytes.Buffer{}) + if err := app.Run([]string{ + "capture", "--harness", "copilot", "--event", "sessionEnd", + "--state-dir", stateDir, "--payload-file", payloadFile, "--no-reconcile", + }); err != nil { + t.Fatalf("capture returned error: %v", err) + } + + db, err := store.Open(stateDir) + if err != nil { + t.Fatalf("Open returned error: %v", err) + } + before, err := db.GetSession(sessionID) + db.Close() + if err != nil { + t.Fatalf("GetSession returned error: %v", err) + } + if before.CaptureState != sessionpkg.CaptureStateProvisional { + t.Fatalf("pre-condition: CaptureState = %q, want provisional", before.CaptureState) + } + + // Stage 2: shutdown event arrives → reconcile upgrades to final. + if err := os.WriteFile(eventsFile, []byte(sessionFinalJSONL), 0o644); err != nil { + t.Fatalf("WriteFile(final) returned error: %v", err) + } + reconcilePayload := writeTempFile(t, "reconcile_payload.json", payload) + + if err := app.Run([]string{ + "reconcile", "--harness", "copilot", "--event", "sessionEnd", + "--state-dir", stateDir, + "--payload-file", reconcilePayload, + "--session-file", eventsFile, + "--session-id", sessionID, + "--attempts", "1", + }); err != nil { + t.Fatalf("reconcile returned error: %v", err) + } + + db2, err := store.Open(stateDir) + if err != nil { + t.Fatalf("Open returned error: %v", err) + } + defer db2.Close() + after, err := db2.GetSession(sessionID) + if err != nil { + t.Fatalf("GetSession returned error: %v", err) + } + if after.CaptureState != sessionpkg.CaptureStateFinal { + t.Errorf("CaptureState = %q, want final", after.CaptureState) + } + if after.TotalInputTokens == nil || *after.TotalInputTokens != 500 { + t.Errorf("TotalInputTokens = %#v, want 500", after.TotalInputTokens) + } +} + +func TestRunReconcileMarksBestEffortWhenExhausted(t *testing.T) { + stateDir := t.TempDir() + copilotHome := t.TempDir() + sessionID := "sess-reconcile-best-effort" + + eventsDir := filepath.Join(copilotHome, "session-state", sessionID) + if err := os.MkdirAll(eventsDir, 0o755); err != nil { + t.Fatalf("MkdirAll returned error: %v", err) + } + eventsFile := filepath.Join(eventsDir, "events.jsonl") + if err := os.WriteFile(eventsFile, []byte(sessionProvisionalJSONL), 0o644); err != nil { + t.Fatalf("WriteFile returned error: %v", err) + } + + payload := fmt.Sprintf(`{"sessionId":%q,"cwd":%q}`, sessionID, t.TempDir()) + payloadFile := writeTempFile(t, "payload.json", payload) + t.Setenv("COPILOT_HOME", copilotHome) + t.Setenv("AROK_COPILOT_SHUTDOWN_RETRY_ATTEMPTS", "1") + + app := New(bytes.NewReader(nil), &bytes.Buffer{}, &bytes.Buffer{}) + if err := app.Run([]string{ + "capture", "--harness", "copilot", "--event", "sessionEnd", + "--state-dir", stateDir, "--payload-file", payloadFile, "--no-reconcile", + }); err != nil { + t.Fatalf("capture returned error: %v", err) + } + + // Reconcile against the still-provisional events file — should exhaust and mark best_effort. + reconcilePayload := writeTempFile(t, "reconcile_payload.json", payload) + err := app.Run([]string{ + "reconcile", "--harness", "copilot", "--event", "sessionEnd", + "--state-dir", stateDir, + "--payload-file", reconcilePayload, + "--session-file", eventsFile, + "--session-id", sessionID, + "--attempts", "1", + }) + if err == nil { + t.Fatal("expected error when reconcile exhausts without final metrics") + } + + db, err := store.Open(stateDir) + if err != nil { + t.Fatalf("Open returned error: %v", err) + } + defer db.Close() + after, err := db.GetSession(sessionID) + if err != nil { + t.Fatalf("GetSession returned error: %v", err) + } + if after.CaptureState != sessionpkg.CaptureStateBestEffort { + t.Errorf("CaptureState = %q, want best_effort", after.CaptureState) + } +} + +func writeTempFile(t *testing.T, name, content string) string { + t.Helper() + path := filepath.Join(t.TempDir(), name) + if err := os.WriteFile(path, []byte(content), 0o644); err != nil { + t.Fatalf("WriteFile(%s) returned error: %v", name, err) + } + return path +} diff --git a/internal/cli/app_vscode.go b/internal/cli/app_vscode.go new file mode 100644 index 0000000..f010a75 --- /dev/null +++ b/internal/cli/app_vscode.go @@ -0,0 +1,280 @@ +package cli + +// VS Code Copilot harness — capture support (scan and per-session Stop events). + +import ( + "encoding/json" + "errors" + "fmt" + "os" + "path/filepath" + "slices" + "strings" + "time" + + "github.com/srbouffard/arok/internal/config" + "github.com/srbouffard/arok/internal/gitmeta" + sessionpkg "github.com/srbouffard/arok/internal/session" + "github.com/srbouffard/arok/internal/store" + "github.com/srbouffard/arok/internal/vscode" +) + +// runCaptureVSCode handles: arok capture --harness vscode --event +// +// Supported events: +// - scan : scans the VS Code workspaceStorage directory and imports all sessions +// - Stop : handles a single VS Code Copilot Stop hook payload +func (a *App) runCaptureVSCode(eventName, stateDirOverride, payloadFile string) error { + if eventName == "" { + return errors.New("missing --event") + } + + stateDir, err := config.ResolveStateDir(stateDirOverride) + if err != nil { + return err + } + if err := config.EnsureLayout(stateDir); err != nil { + return err + } + + db, err := store.Open(stateDir) + if err != nil { + return err + } + defer db.Close() + + if strings.EqualFold(eventName, "scan") { + return a.captureVSCodeScan(stateDir, db) + } + + payloadRaw, err := readPayload(a.stdin, payloadFile) + if err != nil { + return err + } + + payload, ok := parseVSCodeStopPayload(payloadRaw) + if !ok { + if len(strings.TrimSpace(string(payloadRaw))) > 0 { + _ = appendLog(config.IngestLogPath(stateDir), fmt.Sprintf("%s VS Code capture skipped: invalid Stop payload\n", time.Now().UTC().Format(time.RFC3339Nano))) + } + return nil + } + + sessionPath := deriveVSCodeChatSessionPath(payload.TranscriptPath) + if sessionPath == "" { + _ = appendLog(config.IngestLogPath(stateDir), fmt.Sprintf("%s VS Code capture skipped for %s: unable to derive chatSessions path\n", time.Now().UTC().Format(time.RFC3339Nano), payload.SessionID)) + return nil + } + + sessionData, err := vscode.ReadChatSession(sessionPath) + if err != nil { + _ = appendLog(config.IngestLogPath(stateDir), fmt.Sprintf("%s VS Code capture skipped for %s: %v\n", time.Now().UTC().Format(time.RFC3339Nano), payload.SessionID, err)) + return nil + } + if payload.SessionID != "" { + sessionData.SessionID = payload.SessionID + } + + summary := buildVSCodeSummary(sessionData, stateDir, eventName, sessionPath, payload.TranscriptPath) + if summary.SessionID == "" { + _ = appendLog(config.IngestLogPath(stateDir), fmt.Sprintf("%s VS Code capture skipped: missing session_id\n", time.Now().UTC().Format(time.RFC3339Nano))) + return nil + } + + if err := upsertVSCodeSummary(db, summary); err != nil { + _ = appendLog(config.IngestLogPath(stateDir), fmt.Sprintf("%s VS Code database write failed for %s: %v\n", time.Now().UTC().Format(time.RFC3339Nano), summary.SessionID, err)) + return nil + } + return nil +} + +func (a *App) captureVSCodeScan(stateDir string, db *store.Store) error { + userDataDir := vscode.DefaultUserDataDir() + if userDataDir == "" { + _ = appendLog(config.IngestLogPath(stateDir), fmt.Sprintf("%s VS Code scan skipped: user data directory not determinable\n", time.Now().UTC().Format(time.RFC3339Nano))) + return nil + } + + sessions, err := vscode.ScanSessions(userDataDir) + if err != nil { + _ = appendLog(config.IngestLogPath(stateDir), fmt.Sprintf("%s VS Code scan skipped: %v\n", time.Now().UTC().Format(time.RFC3339Nano), err)) + return nil + } + + for _, sessionData := range sessions { + if !shouldImportScannedVSCodeSession(sessionData) { + continue + } + summary := buildVSCodeSummary(sessionData, stateDir, "scan", "", "") + if summary.SessionID == "" { + continue + } + if err := upsertVSCodeSummary(db, summary); err != nil { + _ = appendLog(config.IngestLogPath(stateDir), fmt.Sprintf("%s VS Code scan write failed for %s: %v\n", time.Now().UTC().Format(time.RFC3339Nano), summary.SessionID, err)) + } + } + return nil +} + +type vscodeStopPayload struct { + HookEventName string `json:"hook_event_name"` + SessionID string `json:"session_id"` + TranscriptPath string `json:"transcript_path"` + CWD string `json:"cwd"` +} + +func parseVSCodeStopPayload(raw []byte) (vscodeStopPayload, bool) { + if len(strings.TrimSpace(string(raw))) == 0 { + return vscodeStopPayload{}, false + } + var payload vscodeStopPayload + if err := json.Unmarshal(raw, &payload); err != nil { + return vscodeStopPayload{}, false + } + if payload.SessionID == "" || payload.TranscriptPath == "" { + return vscodeStopPayload{}, false + } + return payload, true +} + +func deriveVSCodeChatSessionPath(transcriptPath string) string { + if strings.TrimSpace(transcriptPath) == "" { + return "" + } + sessionFile := filepath.Base(transcriptPath) + storageDir := filepath.Dir(filepath.Dir(filepath.Dir(transcriptPath))) + if sessionFile == "." || sessionFile == string(filepath.Separator) { + return "" + } + return filepath.Join(storageDir, "chatSessions", sessionFile) +} + +func buildVSCodeSummary(sessionData vscode.Session, stateDir, eventName, eventLogPath, transcriptPath string) sessionpkg.SessionSummary { + hostName, _ := os.Hostname() + git := gitmeta.Inspect(sessionData.WorkspaceFolder) + + var ( + totalOutputTokens int64 + totalInputTokens int64 + modelStats = map[string]int64{} + modelCounts = map[string]int64{} + endedAt = sessionData.CreationDate + ) + for _, req := range sessionData.Requests { + totalOutputTokens += req.CompletionTokens + totalInputTokens += req.PromptTokens + model := req.ModelID + if model == "" { + model = "unknown" + } + modelStats[model] += req.CompletionTokens + modelCounts[model]++ + if req.Timestamp.After(endedAt) { + endedAt = req.Timestamp + } + } + + modelNames := make([]string, 0, len(modelStats)) + for model := range modelStats { + modelNames = append(modelNames, model) + } + slices.Sort(modelNames) + + models := make([]sessionpkg.ModelUsage, 0, len(modelNames)) + for _, model := range modelNames { + tokens := modelStats[model] + count := modelCounts[model] + models = append(models, sessionpkg.ModelUsage{ + Model: model, + AssistantMessageCount: count, + AssistantOutputTokens: tokens, + OutputTokens: sessionpkg.PtrInt64(tokens), + RequestCount: sessionpkg.PtrInt64(count), + }) + } + + startedAt := "" + if !sessionData.CreationDate.IsZero() { + startedAt = sessionData.CreationDate.UTC().Format(time.RFC3339) + } + endedAtRaw := "" + if !endedAt.IsZero() { + endedAtRaw = endedAt.UTC().Format(time.RFC3339) + } + + notes := []string{ + "VS Code Copilot usage is summarized from the local chatSessions JSONL transaction log.", + "Token counts are sourced from result.metadata (primary) or completionTokens/usage fields (fallback).", + } + if sessionData.WorkspaceFolder == "" { + notes = append(notes, "Workspace metadata is unavailable for missing or remote VS Code workspaces.") + } + + var totalInputTokensPtr *int64 + if totalInputTokens > 0 { + totalInputTokensPtr = sessionpkg.PtrInt64(totalInputTokens) + } + + return sessionpkg.SessionSummary{ + SchemaVersion: 1, + Source: "vscode", + Harness: sessionpkg.HarnessVSCodeCopilot, + CollectedAt: time.Now().UTC().Format(time.RFC3339Nano), + SessionID: sessionData.SessionID, + EventName: eventName, + CaptureState: sessionpkg.CaptureStateFinal, + UsageSource: "chatSessions.completionTokens", + TranscriptPath: transcriptPath, + EventLogPath: eventLogPath, + StateDir: stateDir, + CWD: sessionData.WorkspaceFolder, + RepoRoot: git.RepoRoot, + WorktreeRoot: git.WorktreeRoot, + GitCommonDir: git.GitCommonDir, + RepoRemote: git.RepoRemote, + RepoBranch: git.RepoBranch, + RepoHead: git.RepoHead, + HostName: hostName, + StartedAt: startedAt, + EndedAt: endedAtRaw, + InteractionCount: int64(len(sessionData.Requests)), + AssistantMessageCount: int64(len(sessionData.Requests)), + AssistantOutputTokens: totalOutputTokens, + TotalInputTokens: totalInputTokensPtr, + TotalOutputTokens: sessionpkg.PtrInt64(totalOutputTokens), + Models: models, + Notes: notes, + } +} + +func shouldImportScannedVSCodeSession(sessionData vscode.Session) bool { + if len(sessionData.Requests) == 0 { + return true + } + for _, req := range sessionData.Requests { + if req.CompletionTokens > 0 { + return true + } + } + return false +} + +func upsertVSCodeSummary(db *store.Store, summary sessionpkg.SessionSummary) error { + existing, err := db.GetSession(summary.SessionID) + if err != nil && !errors.Is(err, store.ErrSessionNotFound) { + return err + } + if err == nil && existing.CaptureState == sessionpkg.CaptureStateFinal && + derefInt64(existing.TotalOutputTokens) == derefInt64(summary.TotalOutputTokens) && + derefInt64(existing.TotalInputTokens) == derefInt64(summary.TotalInputTokens) { + return nil + } + return db.UpsertSession(summary) +} + +func derefInt64(value *int64) int64 { + if value == nil { + return 0 + } + return *value +} diff --git a/internal/copilot/testdata/session_final.jsonl b/internal/copilot/testdata/session_final.jsonl new file mode 100644 index 0000000..7afa8ba --- /dev/null +++ b/internal/copilot/testdata/session_final.jsonl @@ -0,0 +1,4 @@ +{"timestamp":"2026-01-01T00:00:00Z","type":"assistant.message","data":{"model":"claude-sonnet","outputTokens":150,"interactionId":"i-1","toolRequests":[{}]}} +{"timestamp":"2026-01-01T00:00:01Z","type":"tool.execution_complete","data":{"success":true}} +{"timestamp":"2026-01-01T00:00:02Z","type":"subagent.completed","agentId":"agent-1","data":{"toolCallId":"tc-1","agentName":"explore","agentDisplayName":"Explore","model":"claude-haiku","totalToolCalls":4,"totalTokens":800,"durationMs":2500}} +{"timestamp":"2026-01-01T00:00:03Z","type":"session.shutdown","data":{"modelMetrics":{"claude-sonnet":{"usage":{"inputTokens":500,"outputTokens":150,"cacheReadTokens":200,"cacheWriteTokens":50},"requests":{"count":1}},"claude-haiku":{"usage":{"inputTokens":300,"outputTokens":100},"requests":{"count":4}}}}} diff --git a/internal/copilot/testdata/session_provisional.jsonl b/internal/copilot/testdata/session_provisional.jsonl new file mode 100644 index 0000000..e162cc7 --- /dev/null +++ b/internal/copilot/testdata/session_provisional.jsonl @@ -0,0 +1,3 @@ +{"timestamp":"2026-01-01T00:00:00Z","type":"assistant.message","data":{"model":"claude-sonnet","outputTokens":150,"interactionId":"i-1","toolRequests":[{}]}} +{"timestamp":"2026-01-01T00:00:01Z","type":"tool.execution_complete","data":{"success":true}} +{"timestamp":"2026-01-01T00:00:02Z","type":"assistant.message","data":{"model":"claude-sonnet","outputTokens":80,"interactionId":"i-2"}}