MUL-4921: fix(codex): report resumed usage as task delta

server/pkg/agent/codex.go

@@ -1212,7 +1212,7 @@ func (b *codexBackend) executeOnce(ctx context.Context, prompt string, opts Exec
 		// there rather than in the shared ~/.codex/sessions (MUL-4424).
 		if u.InputTokens == 0 && u.OutputTokens == 0 {
 			taskCodexHome := strings.TrimSpace(b.cfg.Env["CODEX_HOME"])
-			if scanned := scanCodexSessionUsage(startTime, taskCodexHome); scanned != nil {
+			if scanned := scanCodexSessionUsage(startTime, taskCodexHome, threadID, resumed); scanned != nil {
 				u = scanned.usage
 				if scanned.model != "" && opts.Model == "" {
 					opts.Model = scanned.model
@@ -2255,46 +2255,140 @@ type codexSessionUsage struct {
 	model string
 }
 
-// scanCodexSessionUsage scans Codex session JSONL files written after startTime
-// to extract token usage. Codex writes token_count events to
-// $CODEX_HOME/sessions/YYYY/MM/DD/*.jsonl. codexHome is the backend's per-task
-// CODEX_HOME; sessions are isolated there rather than in the shared
-// ~/.codex/sessions (MUL-4424), so usage must be read from it.
-func scanCodexSessionUsage(startTime time.Time, codexHome string) *codexSessionUsage {
+// scanCodexSessionUsage extracts usage for threadID from its Codex rollout.
+// Codex 0.144.4 embeds the app-server thread ID in both the rollout filename and
+// the first session_meta item. Binding the scan to that ID prevents a
+// concurrently-created rollout (for example a Codex subagent) from being billed
+// to this task. A resumed rollout keeps its original date directory, so both flat
+// and YYYY/MM/DD layouts are searched.
+func scanCodexSessionUsage(startTime time.Time, codexHome, threadID string, resumed bool) *codexSessionUsage {
 	root := codexSessionRoot(codexHome)
-	if root == "" {
+	if root == "" || strings.TrimSpace(threadID) == "" {
 		return nil
 	}
 
-	// Look in today's session directory.
-	dateDir := filepath.Join(root,
-		fmt.Sprintf("%04d", startTime.Year()),
-		fmt.Sprintf("%02d", int(startTime.Month())),
-		fmt.Sprintf("%02d", startTime.Day()),
-	)
+	type candidate struct {
+		path    string
+		modTime time.Time
+	}
+	var files []candidate
+	for _, path := range findCodexSessionRollouts(root, threadID) {
+		info, err := os.Stat(path)
+		if err != nil || info.ModTime().Before(startTime) {
+			continue
+		}
+		files = append(files, candidate{path: path, modTime: info.ModTime()})
+	}
+	if len(files) == 0 {
+		return nil
+	}
+	sort.Slice(files, func(i, j int) bool {
+		if files[i].modTime.Equal(files[j].modTime) {
+			return files[i].path < files[j].path
+		}
+		return files[i].modTime.Before(files[j].modTime)
+	})
 
-	files, err := filepath.Glob(filepath.Join(dateDir, "*.jsonl"))
-	if err != nil || len(files) == 0 {
+	// Multiple paths for one thread can transiently exist during layout migration.
+	// They have the same owner, so prefer the latest deterministically without ever
+	// crossing into a different thread's rollout.
+	result := parseCodexSessionFileSince(files[len(files)-1].path, startTime, resumed)
+	if result == nil || (result.usage.InputTokens == 0 && result.usage.OutputTokens == 0 &&
+		result.usage.CacheReadTokens == 0 && result.usage.CacheWriteTokens == 0) {
 		return nil
 	}
+	return result
+}
 
-	// Only scan files modified after startTime (this task's session).
-	var result codexSessionUsage
-	for _, f := range files {
-		info, err := os.Stat(f)
-		if err != nil || info.ModTime().Before(startTime) {
+// findCodexSessionRollouts returns uncompressed rollout files owned by threadID
+// in the two layouts supported by Codex 0.14x. filepath.Glob traverses a linked
+// sessions root, unlike WalkDir, which treats a symlink root as a single file and
+// never visits its children. The normal path uses Codex's filename contract and
+// validates session_meta.id when present. If a future Codex version changes the
+// filename, the metadata pass preserves exact ownership instead of silently
+// dropping usage. Filtering after globbing also keeps a provider-supplied thread
+// ID out of the glob expression.
+func findCodexSessionRollouts(root, threadID string) []string {
+	threadID = strings.TrimSpace(threadID)
+	if root == "" || threadID == "" {
+		return nil
+	}
+
+	patterns := []string{
+		filepath.Join(root, "rollout-*.jsonl"),
+		filepath.Join(root, "*", "*", "*", "rollout-*.jsonl"),
+	}
+	seen := make(map[string]bool)
+	var candidates []string
+	for _, pattern := range patterns {
+		paths, err := filepath.Glob(pattern)
+		if err != nil {
 			continue
 		}
-		if u := parseCodexSessionFile(f); u != nil {
-			// Take the last matching file's data (usually there's only one per task).
-			result = *u
+		for _, path := range paths {
+			if seen[path] {
+				continue
+			}
+			seen[path] = true
+			candidates = append(candidates, path)
 		}
 	}
 
-	if result.usage.InputTokens == 0 && result.usage.OutputTokens == 0 {
-		return nil
+	// Fast path for the current Codex contract. Reject a filename match if its
+	// canonical session_meta explicitly names a different thread.
+	suffix := "-" + threadID + ".jsonl"
+	var matches []string
+	for _, path := range candidates {
+		if !strings.HasSuffix(filepath.Base(path), suffix) {
+			continue
+		}
+		if metadataID, ok := readCodexRolloutThreadID(path); ok && metadataID != threadID {
+			continue
+		}
+		matches = append(matches, path)
 	}
-	return &result
+	if len(matches) > 0 {
+		return matches
+	}
+
+	// Compatibility path for a future filename format: the first session_meta
+	// remains the canonical owner of a rollout in Codex's own resume reader.
+	for _, path := range candidates {
+		if metadataID, ok := readCodexRolloutThreadID(path); ok && metadataID == threadID {
+			matches = append(matches, path)
+		}
+	}
+	return matches
+}
+
+// readCodexRolloutThreadID reads only the head of a rollout. Codex writes the
+// canonical session_meta first; the line cap prevents a malformed legacy file
+// without metadata from turning ownership checks into a full multi-GB scan.
+func readCodexRolloutThreadID(path string) (string, bool) {
+	f, err := os.Open(path)
+	if err != nil {
+		return "", false
+	}
+	defer f.Close()
+
+	var evt struct {
+		Type    string `json:"type"`
+		Payload *struct {
+			ID string `json:"id"`
+		} `json:"payload"`
+	}
+	scanner := bufio.NewScanner(f)
+	scanner.Buffer(make([]byte, 0, 64*1024), 1024*1024)
+	for lineCount := 0; lineCount < 64 && scanner.Scan(); lineCount++ {
+		line := scanner.Bytes()
+		if !bytesContainsStr(line, "session_meta") {
+			continue
+		}
+		if err := json.Unmarshal(line, &evt); err == nil && evt.Type == "session_meta" && evt.Payload != nil && evt.Payload.ID != "" {
+			return evt.Payload.ID, true
+		}
+	}
+	return "", false
 }
 
 // codexSessionRoot returns the Codex sessions directory. It prefers the
@@ -2324,42 +2418,52 @@ func codexSessionRoot(codexHome string) string {
 	return ""
 }
 
+type codexRawTokenUsage struct {
+	InputTokens           int64 `json:"input_tokens"`
+	OutputTokens          int64 `json:"output_tokens"`
+	CachedInputTokens     int64 `json:"cached_input_tokens"`
+	CacheReadInputTokens  int64 `json:"cache_read_input_tokens"`
+	ReasoningOutputTokens int64 `json:"reasoning_output_tokens"`
+}
+
 // codexSessionTokenCount represents a token_count event in Codex JSONL.
 type codexSessionTokenCount struct {
-	Type    string `json:"type"`
-	Payload *struct {
+	Timestamp time.Time `json:"timestamp"`
+	Type      string    `json:"type"`
+	Payload   *struct {
 		Type string `json:"type"`
 		Info *struct {
-			TotalTokenUsage *struct {
-				InputTokens           int64 `json:"input_tokens"`
-				OutputTokens          int64 `json:"output_tokens"`
-				CachedInputTokens     int64 `json:"cached_input_tokens"`
-				CacheReadInputTokens  int64 `json:"cache_read_input_tokens"`
-				ReasoningOutputTokens int64 `json:"reasoning_output_tokens"`
-			} `json:"total_token_usage"`
-			LastTokenUsage *struct {
-				InputTokens           int64 `json:"input_tokens"`
-				OutputTokens          int64 `json:"output_tokens"`
-				CachedInputTokens     int64 `json:"cached_input_tokens"`
-				CacheReadInputTokens  int64 `json:"cache_read_input_tokens"`
-				ReasoningOutputTokens int64 `json:"reasoning_output_tokens"`
-			} `json:"last_token_usage"`
-			Model string `json:"model"`
+			TotalTokenUsage *codexRawTokenUsage `json:"total_token_usage"`
+			LastTokenUsage  *codexRawTokenUsage `json:"last_token_usage"`
+			Model           string              `json:"model"`
 		} `json:"info"`
 		Model string `json:"model"`
 	} `json:"payload"`
 }
 
 // parseCodexSessionFile extracts the final token_count from a Codex session file.
 func parseCodexSessionFile(path string) *codexSessionUsage {
+	return parseCodexSessionFileSince(path, time.Time{}, false)
+}
+
+// parseCodexSessionFileSince extracts usage accumulated after startTime. Codex
+// reports total_token_usage cumulatively for the whole resumed session, so the
+// last total before startTime is subtracted from the final total. Timestamp-less
+// events in a resumed rollout are baseline-only until an explicit post-start
+// timestamp establishes the boundary; fresh sessions retain the previous
+// whole-file behavior because every event belongs to the new task.
+func parseCodexSessionFileSince(path string, startTime time.Time, resumed bool) *codexSessionUsage {
 	f, err := os.Open(path)
 	if err != nil {
 		return nil
 	}
 	defer f.Close()
 
 	var result codexSessionUsage
-	found := false
+	var previousTotal, accumulated, finalUsage codexRawTokenUsage
+	previousTotalFound := false
+	finalUsageFound := false
+	afterStartBoundary := false
 
 	scanner := bufio.NewScanner(f)
 	scanner.Buffer(make([]byte, 0, 256*1024), 1024*1024)
@@ -2376,6 +2480,10 @@ func parseCodexSessionFile(path string) *codexSessionUsage {
 		if err := json.Unmarshal(line, &evt); err != nil || evt.Payload == nil {
 			continue
 		}
+		timestampAfterStart := !startTime.IsZero() && !evt.Timestamp.IsZero() && evt.Timestamp.After(startTime)
+		if timestampAfterStart {
+			afterStartBoundary = true
+		}
 
 		// Track model from turn_context events.
 		if evt.Type == "turn_context" && evt.Payload.Model != "" {
@@ -2385,34 +2493,84 @@ func parseCodexSessionFile(path string) *codexSessionUsage {
 
 		// Extract token usage from token_count events.
 		if evt.Payload.Type == "token_count" && evt.Payload.Info != nil {
-			usage := evt.Payload.Info.TotalTokenUsage
-			if usage == nil {
-				usage = evt.Payload.Info.LastTokenUsage
-			}
-			if usage != nil {
-				cachedTokens := usage.CachedInputTokens
-				if cachedTokens == 0 {
-					cachedTokens = usage.CacheReadInputTokens
-				}
-				result.usage = TokenUsage{
-					InputTokens:     codexUncachedInputTokens(usage.InputTokens, cachedTokens),
-					OutputTokens:    usage.OutputTokens + usage.ReasoningOutputTokens,
-					CacheReadTokens: cachedTokens,
-				}
-				if evt.Payload.Info.Model != "" {
-					result.model = evt.Payload.Info.Model
+			afterStart := startTime.IsZero() || timestampAfterStart ||
+				(evt.Timestamp.IsZero() && (!resumed || afterStartBoundary))
+			if usage := evt.Payload.Info.TotalTokenUsage; usage != nil {
+				current := normalizeCodexRawTokenUsage(*usage)
+				if afterStart {
+					delta := current
+					if previousTotalFound {
+						delta = subtractCodexRawTokenUsage(current, previousTotal)
+					}
+					accumulated = addCodexRawTokenUsage(accumulated, delta)
+					finalUsage = accumulated
+					finalUsageFound = true
 				}
-				found = true
+				previousTotal = current
+				previousTotalFound = true
+			} else if usage := evt.Payload.Info.LastTokenUsage; usage != nil && afterStart {
+				// Preserve event order: a later last_token_usage is the same
+				// fallback the old whole-file parser would have selected.
+				finalUsage = normalizeCodexRawTokenUsage(*usage)
+				finalUsageFound = true
+			}
+			if evt.Payload.Info.Model != "" {
+				result.model = evt.Payload.Info.Model
 			}
 		}
 	}
 
-	if !found {
+	if !finalUsageFound {
 		return nil
 	}
+	cachedTokens := finalUsage.CachedInputTokens
+	result.usage = TokenUsage{
+		InputTokens:     codexUncachedInputTokens(finalUsage.InputTokens, cachedTokens),
+		OutputTokens:    finalUsage.OutputTokens + finalUsage.ReasoningOutputTokens,
+		CacheReadTokens: cachedTokens,
+	}
 	return &result
 }
 
+func subtractCodexRawTokenUsage(total, baseline codexRawTokenUsage) codexRawTokenUsage {
+	total = normalizeCodexRawTokenUsage(total)
+	baseline = normalizeCodexRawTokenUsage(baseline)
+	// Guard each counter independently. Treating one field's reset as a reset of
+	// the entire snapshot would re-report still-monotonic fields and recreate the
+	// over-counting this fallback is meant to prevent.
+	return codexRawTokenUsage{
+		InputTokens:           nonNegativeTokenDelta(total.InputTokens, baseline.InputTokens),
+		OutputTokens:          nonNegativeTokenDelta(total.OutputTokens, baseline.OutputTokens),
+		CachedInputTokens:     nonNegativeTokenDelta(total.CachedInputTokens, baseline.CachedInputTokens),
+		ReasoningOutputTokens: nonNegativeTokenDelta(total.ReasoningOutputTokens, baseline.ReasoningOutputTokens),
+	}
+}
+
+func normalizeCodexRawTokenUsage(usage codexRawTokenUsage) codexRawTokenUsage {
+	if usage.CachedInputTokens == 0 {
+		usage.CachedInputTokens = usage.CacheReadInputTokens
+	}
+	usage.CacheReadInputTokens = 0
+	return usage
+}
+
+func addCodexRawTokenUsage(a, b codexRawTokenUsage) codexRawTokenUsage {
+	return codexRawTokenUsage{
+		InputTokens:           a.InputTokens + b.InputTokens,
+		OutputTokens:          a.OutputTokens + b.OutputTokens,
+		CachedInputTokens:     a.CachedInputTokens + b.CachedInputTokens,
+		ReasoningOutputTokens: a.ReasoningOutputTokens + b.ReasoningOutputTokens,
+	}
+}
+
+func nonNegativeTokenDelta(total, baseline int64) int64 {
+	if total < baseline {
+		// A counter reset means the final value already belongs to the new span.
+		return total
+	}
+	return total - baseline
+}
+
 // bytesContainsStr checks if b contains the string s (without allocating).
 func bytesContainsStr(b []byte, s string) bool {
 	return strings.Contains(string(b), s)