From 3ac08d2f98ec57707344ab77e378a5bd9e4159b8 Mon Sep 17 00:00:00 2001 From: Samuel Bouffard Date: Thu, 18 Jun 2026 14:58:56 +0200 Subject: [PATCH 1/5] chore: split harness logic into per-file files and add docs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Extract copilot capture/reconcile/install logic from app.go into app_copilot.go; VS Code logic into app_vscode.go. app.go is now generic dispatch only — the switch in runCapture() is the natural extension point for new harnesses. - Add app_copilot_test.go with integration tests covering the final, provisional, missing-sessionId, and unknown-harness cases. - Add internal/copilot/testdata/ fixtures (session_final.jsonl and session_provisional.jsonl) as canonical examples of both states. - Add docs/adding-a-harness.md: step-by-step guide for contributors including file structure, SessionSummary field reference, reconcile guidance, test patterns, and a checklist. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- docs/adding-a-harness.md | 246 +++++++ internal/cli/app.go | 638 +----------------- internal/cli/app_copilot.go | 393 +++++++++++ internal/cli/app_copilot_test.go | 165 +++++ internal/cli/app_vscode.go | 280 ++++++++ internal/copilot/testdata/session_final.jsonl | 4 + .../testdata/session_provisional.jsonl | 3 + 7 files changed, 1103 insertions(+), 626 deletions(-) create mode 100644 docs/adding-a-harness.md create mode 100644 internal/cli/app_copilot.go create mode 100644 internal/cli/app_copilot_test.go create mode 100644 internal/cli/app_vscode.go create mode 100644 internal/copilot/testdata/session_final.jsonl create mode 100644 internal/copilot/testdata/session_provisional.jsonl diff --git a/docs/adding-a-harness.md b/docs/adding-a-harness.md new file mode 100644 index 0000000..6dc95b3 --- /dev/null +++ b/docs/adding-a-harness.md @@ -0,0 +1,246 @@ +# 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 walks you through +adding support for a new harness. + +## Overview + +The capture pipeline for every harness looks the same: + +``` +agent tool fires hook + → runs: arok capture --harness --event + → parses the hook payload + → reads session data from the harness-specific source + → builds a session.SessionSummary + → stores it in the arok SQLite database +``` + +Adding a harness means implementing steps 2–4 for your tool. + +## File structure + +The codebase organises harness-specific code into two places: + +``` +internal// ← payload parsing and session summarizing logic +internal/cli/app_.go ← CLI plumbing: runCapture(), helpers +``` + +Look at the existing copilot harness for reference: + +``` +internal/copilot/copilot.go ← Summarize(), ParsePayload(), ResolveSessionFile() +internal/cli/app_copilot.go ← runCaptureCopilot(), runReconcile(), etc. +internal/cli/app_copilot_test.go ← integration tests +internal/copilot/testdata/ ← fixture files for unit tests +``` + +## Step-by-step: adding "myharness" + +### 1. Create the parsing package + +Create `internal/myharness/myharness.go`. Its job is to convert a raw hook payload +and whatever session log/API the tool provides into a `session.SessionSummary`. + +```go +package myharness + +import ( + sessionpkg "github.com/srbouffard/arok/internal/session" + "time" +) + +// Payload holds the data delivered by the hook runner. +type Payload struct { + SessionID string `json:"session_id"` + CWD string `json:"cwd"` + // ... add fields from your tool's hook payload +} + +// Summarize converts payload and session data into a SessionSummary. +func Summarize(eventName, stateDir string, p Payload) (sessionpkg.SessionSummary, error) { + // Read session data from wherever your tool stores it. + // Build and return a SessionSummary. + return sessionpkg.SessionSummary{ + SchemaVersion: 1, + Source: "myharness", + Harness: "myharness", // lowercase kebab-case identifier + CollectedAt: time.Now().UTC().Format(time.RFC3339Nano), + SessionID: p.SessionID, + EventName: eventName, + CaptureState: sessionpkg.CaptureStateFinal, + // ... populate remaining fields + }, nil +} +``` + +Key `SessionSummary` fields to fill in: + +| Field | Description | +|-------|-------------| +| `Harness` | Lowercase kebab-case name stored in the database (e.g. `"myharness"`) | +| `SessionID` | Unique session identifier from the hook payload | +| `CaptureState` | `"final"` if you have complete data now; `"provisional"` if you need a reconcile pass | +| `TotalInputTokens` / `TotalOutputTokens` | `*int64` — use `session.PtrInt64(n)` | +| `Models` | Per-model breakdown as `[]session.ModelUsage` | +| `CWD` / `RepoRoot` / `RepoBranch` | Use `gitmeta.Inspect(cwd)` to fill these from the working directory | + +### 2. Create the CLI plumbing + +Create `internal/cli/app_myharness.go`: + +```go +package cli + +import ( + "errors" + "time" + + "github.com/srbouffard/arok/internal/config" + "github.com/srbouffard/arok/internal/myharness" + "github.com/srbouffard/arok/internal/store" +) + +func (a *App) runCaptureMyHarness(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 + } + + payloadRaw, err := readPayload(a.stdin, payloadFile) + if err != nil { + return err + } + + // Parse the hook payload. + var p myharness.Payload + if err := json.Unmarshal(payloadRaw, &p); err != nil { + return err + } + + summary, err := myharness.Summarize(eventName, stateDir, p) + if err != nil { + _ = appendLog(config.IngestLogPath(stateDir), fmt.Sprintf("%s myharness capture failed: %v\n", time.Now().UTC().Format(time.RFC3339Nano), err)) + return err + } + + db, err := store.Open(stateDir) + if err != nil { + return err + } + defer db.Close() + + return db.UpsertSession(summary) +} +``` + +### 3. Wire it into the capture dispatcher + +In `internal/cli/app.go`, add your harness to the `runCapture` switch: + +```go +switch harness { +case "copilot": + return a.runCaptureCopilot(eventName, stateDirOverride, payloadFile, noReconcile) +case "vscode": + return a.runCaptureVSCode(eventName, stateDirOverride, payloadFile) +case "myharness": // ← add this + return a.runCaptureMyHarness(eventName, stateDirOverride, payloadFile) +default: + return fmt.Errorf("unsupported harness %q", harness) +} +``` + +### 4. Wire it into the hook config (optional) + +If your tool uses a JSON hook config file (like Copilot CLI does), add an +`arok install myharness` command by following the same pattern as +`runInstallCopilot` in `app_copilot.go` and `InstallCopilot` in +`internal/install/copilot.go`. + +Then add a case in `runInstall` in `app.go`: + +```go +case "myharness": + return a.runInstallMyHarness(args[1:]) +``` + +### 5. Update the usage string + +In `printRootUsage()` in `app.go`, add your harness to the capture line: + +```go +fmt.Fprintf(a.stdout, "... capture --harness [copilot|vscode|myharness] --event ...") +``` + +## Reconcile (only needed for two-phase capture) + +The copilot harness needs a reconcile pass because `session.shutdown.modelMetrics` +arrives asynchronously after the hook fires. Most harnesses can produce a final +`SessionSummary` immediately and do not need this. + +If your harness fires a hook before all metrics are available, return +`CaptureState: sessionpkg.CaptureStateProvisional` from `Summarize()` and then +spawn a background `arok reconcile --harness myharness` process. See +`spawnDetachedReconcile` and `runReconcile` in `app_copilot.go` for the pattern. + +## Writing tests + +Add two test files: + +**Unit tests** — `internal/myharness/myharness_test.go` + +Test `Summarize()` in isolation using fixture JSONL files in +`internal/myharness/testdata/`. See `internal/copilot/copilot_test.go` for +examples. + +**Integration tests** — `internal/cli/app_myharness_test.go` + +Test the full `arok capture --harness myharness` flow through `App.Run()`. +See `internal/cli/app_copilot_test.go` for the pattern: + +```go +func TestRunCaptureMyHarnessStoresSession(t *testing.T) { + stateDir := t.TempDir() + payloadFile := writeTempFile(t, "payload.json", `{"session_id":"sess-1","cwd":"/tmp"}`) + + app := New(bytes.NewReader(nil), &bytes.Buffer{}, &bytes.Buffer{}) + if err := app.Run([]string{ + "capture", "--harness", "myharness", "--event", "sessionEnd", + "--state-dir", stateDir, + "--payload-file", payloadFile, + }); err != nil { + t.Fatalf("Run returned error: %v", err) + } + + db, _ := store.Open(stateDir) + defer db.Close() + summary, err := db.GetSession("sess-1") + if err != nil { + t.Fatalf("GetSession: %v", err) + } + if summary.Harness != "myharness" { + t.Errorf("Harness = %q, want myharness", summary.Harness) + } + // ... assert token counts, capture state, etc. +} +``` + +## Checklist + +- [ ] `internal/myharness/myharness.go` — `Summarize()` returns a valid `SessionSummary` +- [ ] `internal/cli/app_myharness.go` — `runCaptureMyHarness()` method +- [ ] `internal/cli/app.go` — case added to `runCapture()` switch +- [ ] `internal/myharness/myharness_test.go` — unit tests for `Summarize()` +- [ ] `internal/myharness/testdata/` — fixture JSONL files +- [ ] `internal/cli/app_myharness_test.go` — integration test 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..8f27ca4 --- /dev/null +++ b/internal/cli/app_copilot.go @@ -0,0 +1,393 @@ +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 + } + 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 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..a71ecee --- /dev/null +++ b/internal/cli/app_copilot_test.go @@ -0,0 +1,165 @@ +package cli + +import ( + "bytes" + "fmt" + "os" + "path/filepath" + "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 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"}} From dfc97ed0286e7f44dc44f4720982827a5edc2b1d Mon Sep 17 00:00:00 2001 From: Samuel Bouffard Date: Thu, 18 Jun 2026 16:14:27 +0200 Subject: [PATCH 2/5] test: add reconcile and install coverage; fix docs code example MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Add TestRunReconcileUpgradesProvisionalToFinal: verifies the full provisional→final upgrade path through App.Run(reconcile). - Add TestRunReconcileMarksBestEffortWhenExhausted: verifies that a reconcile that exhausts without finding shutdown metrics marks the session as best_effort and returns an error. - Add TestRunInstallCopilotPrintConfig: verifies --print-config emits valid JSON containing the sessionEnd hook and binary path. - Fix missing imports (encoding/json, fmt, time) in the app_myharness.go code example in docs/adding-a-harness.md. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- docs/adding-a-harness.md | 3 +- internal/cli/app_copilot_test.go | 157 +++++++++++++++++++++++++++++++ 2 files changed, 159 insertions(+), 1 deletion(-) diff --git a/docs/adding-a-harness.md b/docs/adding-a-harness.md index 6dc95b3..787821d 100644 --- a/docs/adding-a-harness.md +++ b/docs/adding-a-harness.md @@ -95,7 +95,9 @@ Create `internal/cli/app_myharness.go`: package cli import ( + "encoding/json" "errors" + "fmt" "time" "github.com/srbouffard/arok/internal/config" @@ -121,7 +123,6 @@ func (a *App) runCaptureMyHarness(eventName, stateDirOverride, payloadFile strin return err } - // Parse the hook payload. var p myharness.Payload if err := json.Unmarshal(payloadRaw, &p); err != nil { return err diff --git a/internal/cli/app_copilot_test.go b/internal/cli/app_copilot_test.go index a71ecee..e1909e7 100644 --- a/internal/cli/app_copilot_test.go +++ b/internal/cli/app_copilot_test.go @@ -2,9 +2,11 @@ package cli import ( "bytes" + "encoding/json" "fmt" "os" "path/filepath" + "strings" "testing" sessionpkg "github.com/srbouffard/arok/internal/session" @@ -155,6 +157,161 @@ func TestRunCaptureCopilotUnknownHarness(t *testing.T) { } } +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) From 96a82825a5de4bc581bdeb1bd3e9eb9a006e8afb Mon Sep 17 00:00:00 2001 From: Samuel Bouffard Date: Thu, 18 Jun 2026 16:21:03 +0200 Subject: [PATCH 3/5] fix: handle logFile.Close() error explicitly in spawnDetachedReconcile defer logFile.Close() on a writable fd without checking the error is a CodeQL-flagged pattern. The child process writes directly to the kernel fd so there is no Go-side buffering to lose, but make the intentional discard explicit with defer func() { _ = logFile.Close() }() to satisfy static analysis and document the reasoning. Pre-existing issue surfaced by CodeQL on this PR's diff. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- internal/cli/app_copilot.go | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/internal/cli/app_copilot.go b/internal/cli/app_copilot.go index 8f27ca4..d69eb04 100644 --- a/internal/cli/app_copilot.go +++ b/internal/cli/app_copilot.go @@ -320,7 +320,10 @@ func spawnDetachedReconcile(stateDir string, payloadRaw []byte, sessionID, event if err != nil { return err } - defer logFile.Close() + // 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, From a0978126719fcf413cf539d315983c66fb070d93 Mon Sep 17 00:00:00 2001 From: Samuel Bouffard Date: Thu, 18 Jun 2026 16:26:42 +0200 Subject: [PATCH 4/5] docs: rewrite harness guide as concept-first, remove inline code Remove boilerplate code examples in favour of directing contributors to the existing harness implementations as the canonical reference. Add explanations of the two key challenges a harness author will encounter: - Hook payloads are thin; the real data lives in a separate log or API that must be located via session ID and CWD. - Final metrics may arrive after the hook exits; use provisional capture + async reconcile rather than blocking the hook runner. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- docs/adding-a-harness.md | 298 ++++++++++++--------------------------- 1 file changed, 91 insertions(+), 207 deletions(-) diff --git a/docs/adding-a-harness.md b/docs/adding-a-harness.md index 787821d..4647611 100644 --- a/docs/adding-a-harness.md +++ b/docs/adding-a-harness.md @@ -1,247 +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 walks you through -adding support for a new harness. +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. -## Overview +## How capture works -The capture pipeline for every harness looks the same: +When an agent session ends, the harness fires a hook that runs: ``` -agent tool fires hook - → runs: arok capture --harness --event - → parses the hook payload - → reads session data from the harness-specific source - → builds a session.SessionSummary - → stores it in the arok SQLite database +arok capture --harness --event ``` -Adding a harness means implementing steps 2–4 for your tool. +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 -The codebase organises harness-specific code into two places: +Each harness lives in two places: ``` -internal// ← payload parsing and session summarizing logic -internal/cli/app_.go ← CLI plumbing: runCapture(), helpers +internal// ← payload parsing and session summarizing logic +internal/cli/app_.go ← CLI plumbing: runCapture() and helpers ``` -Look at the existing copilot harness for reference: +**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(), etc. -internal/cli/app_copilot_test.go ← integration tests -internal/copilot/testdata/ ← fixture files for unit tests +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 ``` -## Step-by-step: adding "myharness" - -### 1. Create the parsing package - -Create `internal/myharness/myharness.go`. Its job is to convert a raw hook payload -and whatever session log/API the tool provides into a `session.SessionSummary`. - -```go -package myharness - -import ( - sessionpkg "github.com/srbouffard/arok/internal/session" - "time" -) - -// Payload holds the data delivered by the hook runner. -type Payload struct { - SessionID string `json:"session_id"` - CWD string `json:"cwd"` - // ... add fields from your tool's hook payload -} - -// Summarize converts payload and session data into a SessionSummary. -func Summarize(eventName, stateDir string, p Payload) (sessionpkg.SessionSummary, error) { - // Read session data from wherever your tool stores it. - // Build and return a SessionSummary. - return sessionpkg.SessionSummary{ - SchemaVersion: 1, - Source: "myharness", - Harness: "myharness", // lowercase kebab-case identifier - CollectedAt: time.Now().UTC().Format(time.RFC3339Nano), - SessionID: p.SessionID, - EventName: eventName, - CaptureState: sessionpkg.CaptureStateFinal, - // ... populate remaining fields - }, nil -} -``` +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 `SessionSummary` fields to fill in: - -| Field | Description | -|-------|-------------| -| `Harness` | Lowercase kebab-case name stored in the database (e.g. `"myharness"`) | -| `SessionID` | Unique session identifier from the hook payload | -| `CaptureState` | `"final"` if you have complete data now; `"provisional"` if you need a reconcile pass | -| `TotalInputTokens` / `TotalOutputTokens` | `*int64` — use `session.PtrInt64(n)` | -| `Models` | Per-model breakdown as `[]session.ModelUsage` | -| `CWD` / `RepoRoot` / `RepoBranch` | Use `gitmeta.Inspect(cwd)` to fill these from the working directory | - -### 2. Create the CLI plumbing - -Create `internal/cli/app_myharness.go`: - -```go -package cli - -import ( - "encoding/json" - "errors" - "fmt" - "time" - - "github.com/srbouffard/arok/internal/config" - "github.com/srbouffard/arok/internal/myharness" - "github.com/srbouffard/arok/internal/store" -) - -func (a *App) runCaptureMyHarness(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 - } - - payloadRaw, err := readPayload(a.stdin, payloadFile) - if err != nil { - return err - } - - var p myharness.Payload - if err := json.Unmarshal(payloadRaw, &p); err != nil { - return err - } - - summary, err := myharness.Summarize(eventName, stateDir, p) - if err != nil { - _ = appendLog(config.IngestLogPath(stateDir), fmt.Sprintf("%s myharness capture failed: %v\n", time.Now().UTC().Format(time.RFC3339Nano), err)) - return err - } - - db, err := store.Open(stateDir) - if err != nil { - return err - } - defer db.Close() - - return db.UpsertSession(summary) -} -``` +## Key concepts -### 3. Wire it into the capture dispatcher - -In `internal/cli/app.go`, add your harness to the `runCapture` switch: - -```go -switch harness { -case "copilot": - return a.runCaptureCopilot(eventName, stateDirOverride, payloadFile, noReconcile) -case "vscode": - return a.runCaptureVSCode(eventName, stateDirOverride, payloadFile) -case "myharness": // ← add this - return a.runCaptureMyHarness(eventName, stateDirOverride, payloadFile) -default: - return fmt.Errorf("unsupported harness %q", harness) -} -``` +### The hook payload may not contain everything you need -### 4. Wire it into the hook config (optional) +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. -If your tool uses a JSON hook config file (like Copilot CLI does), add an -`arok install myharness` command by following the same pattern as -`runInstallCopilot` in `app_copilot.go` and `InstallCopilot` in -`internal/install/copilot.go`. +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. -Then add a case in `runInstall` in `app.go`: +### Final metrics may not be available when the hook fires -```go -case "myharness": - return a.runInstallMyHarness(args[1:]) -``` +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. -### 5. Update the usage string +**Do not block the hook.** The hook runner expects a fast exit. Instead: -In `printRootUsage()` in `app.go`, add your harness to the capture line: +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. -```go -fmt.Fprintf(a.stdout, "... capture --harness [copilot|vscode|myharness] --event ...") -``` +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. -## Reconcile (only needed for two-phase capture) +## Adding your harness: the steps -The copilot harness needs a reconcile pass because `session.shutdown.modelMetrics` -arrives asynchronously after the hook fires. Most harnesses can produce a final -`SessionSummary` immediately and do not need this. +1. **Create `internal//`** — implement `Summarize()`, which takes the parsed + payload and returns a `session.SessionSummary`. Model it on + `internal/copilot/copilot.go`. -If your harness fires a hook before all metrics are available, return -`CaptureState: sessionpkg.CaptureStateProvisional` from `Summarize()` and then -spawn a background `arok reconcile --harness myharness` process. See -`spawnDetachedReconcile` and `runReconcile` in `app_copilot.go` for the pattern. +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 -Add two test files: - -**Unit tests** — `internal/myharness/myharness_test.go` - -Test `Summarize()` in isolation using fixture JSONL files in -`internal/myharness/testdata/`. See `internal/copilot/copilot_test.go` for -examples. - -**Integration tests** — `internal/cli/app_myharness_test.go` - -Test the full `arok capture --harness myharness` flow through `App.Run()`. -See `internal/cli/app_copilot_test.go` for the pattern: - -```go -func TestRunCaptureMyHarnessStoresSession(t *testing.T) { - stateDir := t.TempDir() - payloadFile := writeTempFile(t, "payload.json", `{"session_id":"sess-1","cwd":"/tmp"}`) - - app := New(bytes.NewReader(nil), &bytes.Buffer{}, &bytes.Buffer{}) - if err := app.Run([]string{ - "capture", "--harness", "myharness", "--event", "sessionEnd", - "--state-dir", stateDir, - "--payload-file", payloadFile, - }); err != nil { - t.Fatalf("Run returned error: %v", err) - } - - db, _ := store.Open(stateDir) - defer db.Close() - summary, err := db.GetSession("sess-1") - if err != nil { - t.Fatalf("GetSession: %v", err) - } - if summary.Harness != "myharness" { - t.Errorf("Harness = %q, want myharness", summary.Harness) - } - // ... assert token counts, capture state, etc. -} -``` +**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/myharness/myharness.go` — `Summarize()` returns a valid `SessionSummary` -- [ ] `internal/cli/app_myharness.go` — `runCaptureMyHarness()` method -- [ ] `internal/cli/app.go` — case added to `runCapture()` switch -- [ ] `internal/myharness/myharness_test.go` — unit tests for `Summarize()` -- [ ] `internal/myharness/testdata/` — fixture JSONL files -- [ ] `internal/cli/app_myharness_test.go` — integration test via `App.Run()` +- [ ] `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 From 3393c01c32cfda22a334ee283ae6fc0d9f96f8bd Mon Sep 17 00:00:00 2001 From: Samuel Bouffard Date: Thu, 18 Jun 2026 16:29:04 +0200 Subject: [PATCH 5/5] docs: reference adding-a-harness.md in AGENTS.md Agents tackling a new harness integration will now be directed to the guide and the copilot harness as the canonical reference before writing any code. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- AGENTS.md | 8 ++++++++ 1 file changed, 8 insertions(+) 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 ./...`