Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
107 changes: 61 additions & 46 deletions pkg/cli/token_usage.go
Original file line number Diff line number Diff line change
Expand Up @@ -613,62 +613,77 @@ func sumAICFromUsageJSONLFiles(filePaths []string) (float64, bool, error) {
found := false

for _, filePath := range filePaths {
file, err := os.Open(filepath.Clean(filePath))
fileAIC, fileFound, err := processOneUsageJSONLFile(filePath)
if err != nil {
return 0, false, fmt.Errorf("failed to open usage JSONL file %s: %w", filePath, err)
return 0, false, err
}
totalAIC += fileAIC
if fileFound {
found = true
}
}

scanner := bufio.NewScanner(file)
scanner.Buffer(make([]byte, 0, 64*1024), 1024*1024)
for scanner.Scan() {
line := strings.TrimSpace(scanner.Text())
if line == "" || !strings.HasPrefix(line, "{") {
continue
}
return totalAIC, found, nil
}

var parsed map[string]any
if err := json.Unmarshal([]byte(line), &parsed); err != nil {
continue
}
// processOneUsageJSONLFile reads a single usage JSONL file and returns the total AIC
// accumulated from its records. The file is deferred-closed immediately after open.
func processOneUsageJSONLFile(filePath string) (total float64, found bool, err error) {

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[/tdd] processOneUsageJSONLFile is unexported with no dedicated test — the named-return + defer close-error wiring is subtle enough that a few table-driven tests would give confidence.

💡 Suggested test cases
func TestProcessOneUsageJSONLFile(t *testing.T) {
    t.Run("returns error on missing file", func(t *testing.T) {
        _, _, err := processOneUsageJSONLFile("/nonexistent.jsonl")
        require.Error(t, err)
    })
    t.Run("accumulates AIC from valid lines", func(t *testing.T) {
        f := writeTempJSONL(t, `{"aic":1.5}`)
        total, found, err := processOneUsageJSONLFile(f)
        require.NoError(t, err)
        assert.True(t, found)
        assert.InDelta(t, 1.5, total, 0.001)
    })
}

Covering open-error, partial-parse, and the happy path would also verify that defer file.Close() named-return wiring works correctly under each exit.

file, err := os.Open(filepath.Clean(filePath))
if err != nil {
return 0, false, fmt.Errorf("failed to open usage JSONL file %s: %w", filePath, err)
}
defer func() {

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[/diagnose] The defer close-error logic silently drops a file.Close() error when scanner.Err() has already set err. This matches the original code's behavior (which also dropped closeErr on scanner error), but it's worth a brief comment so future readers understand the intentional priority.

💡 One-liner clarification
defer func() {
    // Close error is only propagated when no other error is already set;
    // scanner errors (set explicitly below) take priority.
    if closeErr := file.Close(); closeErr != nil && err == nil {
        err = fmt.Errorf("failed to close usage JSONL file %s: %w", filePath, closeErr)
    }
}()

The current code is correct — this is purely a documentation suggestion.

if closeErr := file.Close(); closeErr != nil && err == nil {
err = fmt.Errorf("failed to close usage JSONL file %s: %w", filePath, closeErr)
}
}()

usage := extractUsageRecord(parsed["usage"])
explicitAICredits := usageNumericValue(parsed, usage, "ai_credits", "aiCredits")
if explicitAICredits > 0 {
totalAIC += explicitAICredits
found = true
continue
}
explicitAIC := usageNumericValue(parsed, usage, "aic")
if explicitAIC > 0 {
totalAIC += explicitAIC
found = true
continue
}
scanner := bufio.NewScanner(file)
scanner.Buffer(make([]byte, 0, 64*1024), 1024*1024)
for scanner.Scan() {
line := strings.TrimSpace(scanner.Text())
if line == "" || !strings.HasPrefix(line, "{") {
continue
}

computedAIC := computeModelInferenceAIC(
usageStringValue(parsed, usage, "provider"),
usageStringValue(parsed, usage, "model"),
int(usageNumericValue(parsed, usage, "input_tokens", "inputTokens")),
int(usageNumericValue(parsed, usage, "output_tokens", "outputTokens")),
int(usageNumericValue(parsed, usage, "cache_read_tokens", "cacheReadTokens")),
int(usageNumericValue(parsed, usage, "cache_write_tokens", "cacheWriteTokens")),
int(usageNumericValue(parsed, usage, "reasoning_tokens", "reasoningTokens")),
)
if computedAIC > 0 {
totalAIC += computedAIC
found = true
}
var parsed map[string]any
if jsonErr := json.Unmarshal([]byte(line), &parsed); jsonErr != nil {
continue
}
closeErr := file.Close()
if err := scanner.Err(); err != nil {
return 0, false, fmt.Errorf("error reading usage JSONL file %s: %w", filePath, err)

usage := extractUsageRecord(parsed["usage"])
explicitAICredits := usageNumericValue(parsed, usage, "ai_credits", "aiCredits")
if explicitAICredits > 0 {
total += explicitAICredits
found = true
continue
}
if closeErr != nil {
return 0, false, fmt.Errorf("failed to close usage JSONL file %s: %w", filePath, closeErr)
explicitAIC := usageNumericValue(parsed, usage, "aic")
if explicitAIC > 0 {
total += explicitAIC
found = true
continue
}
}

return totalAIC, found, nil
computedAIC := computeModelInferenceAIC(
usageStringValue(parsed, usage, "provider"),
usageStringValue(parsed, usage, "model"),
int(usageNumericValue(parsed, usage, "input_tokens", "inputTokens")),
int(usageNumericValue(parsed, usage, "output_tokens", "outputTokens")),
int(usageNumericValue(parsed, usage, "cache_read_tokens", "cacheReadTokens")),
int(usageNumericValue(parsed, usage, "cache_write_tokens", "cacheWriteTokens")),
int(usageNumericValue(parsed, usage, "reasoning_tokens", "reasoningTokens")),
)
if computedAIC > 0 {
total += computedAIC
found = true
}
}
if scanErr := scanner.Err(); scanErr != nil {
return 0, false, fmt.Errorf("error reading usage JSONL file %s: %w", filePath, scanErr)
}
return total, found, nil
}

// analyzeTokenUsageAICOnly parses token usage inputs and computes only TotalAIC.
Expand Down
4 changes: 2 additions & 2 deletions pkg/parser/remote_fetch.go
Original file line number Diff line number Diff line change
Expand Up @@ -624,8 +624,8 @@ func downloadFileViaGit(ctx context.Context, owner, repo, path, ref, host string

// git archive command: git archive --remote=<repo> <ref> <path>
// #nosec G204 -- repoURL, ref, and path are from workflow import configuration authored by the
// developer; exec.Command with separate args (not shell execution) prevents shell injection.
cmd := exec.Command("git", "archive", "--remote="+repoURL, ref, path)
// developer; exec.CommandContext with separate args (not shell execution) prevents shell injection.
cmd := exec.CommandContext(ctx, "git", "archive", "--remote="+repoURL, ref, path)
archiveOutput, err := cmd.Output()
if err != nil {
// If git archive fails, try with git clone + git show as a fallback
Expand Down
39 changes: 27 additions & 12 deletions pkg/workflow/action_resolver.go
Original file line number Diff line number Diff line change
Expand Up @@ -222,28 +222,43 @@ func (r *ActionResolver) resolveFromGitHub(ctx context.Context, repo, version st
return "", fmt.Errorf("failed to resolve %s@%s: exceeded max tag peel depth %d", repo, version, maxTagPeelDepth)
}
resolverLog.Printf("Detected annotated tag for %s@%s (depth %d, tag object SHA: %s), peeling to underlying object", repo, version, depth, sha)
tagPath := fmt.Sprintf("/repos/%s/git/tags/%s", baseRepo, sha)
// Each peel gets its own fresh 30-second timeout derived from the original
// caller context (ctx), not from callCtx, so we don't accidentally shrink
// the budget for subsequent peels.
peelCtx, peelCancel := context.WithTimeout(ctx, 30*time.Second)
cmd2 := ExecGHContext(peelCtx, "api", tagPath, "--jq", "[.object.sha, .object.type] | @tsv")
ForceGHHostEnv(cmd2, "github.com")
output2, peelErr := cmd2.Output()
peelCancel()
if peelErr != nil {
return "", fmt.Errorf("failed to peel annotated tag %s@%s: %w", repo, version, peelErr)
}
sha, objType, err = ParseTagRefTSV(string(output2))
// the budget for subsequent peels. peelTagObject defers cancel within its
// own function scope to satisfy resource-lifecycle lint rules.
sha, objType, err = peelTagObject(ctx, baseRepo, repo, version, sha)
if err != nil {
return "", fmt.Errorf("failed to parse peeled tag API response for %s@%s: %w", repo, version, err)
return "", err
}
}
resolverLog.Printf("Resolved %s@%s to %s SHA: %s", repo, version, objType, sha)

return sha, nil
}

// peelTagObject resolves a single annotated-tag object to its underlying object by
// querying the GitHub API. It is called iteratively for chained tag objects.
// The timeout context is created and immediately deferred within this function so
// that cancel is always called when the function returns (not deferred inside a loop).
func peelTagObject(ctx context.Context, baseRepo, repo, version, sha string) (newSHA, objType string, err error) {

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[/diagnose] No regression test covers the annotated-tag peel loop, so the peelTagObject extraction has no safety net. A lint-compliance refactor is a low-risk change, but a test verifying that chained tag peeling still returns the correct commit SHA would lock in the behavior and make future refactors safer.

💡 What to test

A table-driven test against a mock ExecGHContext that returns successive sha/tag TSV rows until sha/commit would cover:

  • single annotated tag (depth 1)
  • doubly-chained tag (depth 2)
  • max-depth exceeded → expected error
  • API failure mid-peel → expected wrapped error

This also pins the 30-second-per-peel timeout logic, which is easy to break silently.

tagPath := fmt.Sprintf("/repos/%s/git/tags/%s", baseRepo, sha)
// Each peel gets a fresh 30-second timeout derived from the original caller
// context so we don't accidentally shrink the budget across iterations.
peelCtx, peelCancel := context.WithTimeout(ctx, 30*time.Second)
defer peelCancel()
cmd := ExecGHContext(peelCtx, "api", tagPath, "--jq", "[.object.sha, .object.type] | @tsv")
ForceGHHostEnv(cmd, "github.com")
output, peelErr := cmd.Output()
if peelErr != nil {
return "", "", fmt.Errorf("failed to peel annotated tag %s@%s: %w", repo, version, peelErr)
}
newSHA, objType, err = ParseTagRefTSV(string(output))
if err != nil {
return "", "", fmt.Errorf("failed to parse peeled tag API response for %s@%s: %w", repo, version, err)
}
return newSHA, objType, nil
}

// ResolveGhAwRef resolves a branch, tag, or SHA ref in the github/gh-aw
// repository to its full 40-character commit SHA.
// If ref is already a valid full SHA it is returned unchanged.
Expand Down
14 changes: 9 additions & 5 deletions pkg/workflow/safe_outputs_actions.go
Original file line number Diff line number Diff line change
Expand Up @@ -287,11 +287,15 @@ func fetchRemoteActionYAML(repo, subdir, ref string) (*actionYAMLFile, error) {
apiPath := fmt.Sprintf("/repos/%s/contents/%s?ref=%s", repo, contentPath, ref)
safeOutputActionsLog.Printf("Fetching action YAML from: %s", apiPath)

ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second)
defer cancel()
cmd := ExecGHContext(ctx, "api", apiPath, "--jq", ".content")
output, err := cmd.Output()
cancel()
// Wrap the context creation and command execution in a function literal so that
// defer cancel() runs when the literal returns (not when the outer function returns),
// avoiding a defer-inside-loop resource-lifecycle issue.
output, err := func() ([]byte, error) {

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[/zoom-out] This fix uses an inline anonymous function literal, while action_resolver.go in this same PR extracts a named peelTagObject helper for the same lint issue. A named helper like fetchActionFileContent(apiPath string) ([]byte, error) would be more consistent, and more importantly, directly unit-testable without going through fetchRemoteActionYAML.

💡 Suggested refactor
func fetchActionFileContent(apiPath string) ([]byte, error) {
    ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second)
    defer cancel()
    cmd := ExecGHContext(ctx, "api", apiPath, "--jq", ".content")
    return cmd.Output()
}

Then at the call site: output, err := fetchActionFileContent(apiPath)

A single pattern across the codebase also makes future lint rules easier to verify mechanically.

ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[/diagnose] context.Background() is used here (preserved from before this PR), so the 20-second timeout does not respect upstream cancellation from the caller. If fetchRemoteActionYAML is ever called from a context that can be cancelled (e.g., a request with a shorter deadline), the fetch will still run for the full 20 seconds. Consider threading the parent context through if cancellation propagation becomes desirable.

Not introduced by this PR — just surfaced by the refactor making it more visible.

defer cancel()
cmd := ExecGHContext(ctx, "api", apiPath, "--jq", ".content")
return cmd.Output()
}()
if err != nil {
safeOutputActionsLog.Printf("Failed to fetch %s from %s@%s: %v", filename, repo, ref, err)
continue
Expand Down
Loading