Skip to content

Commit 7962a28

Browse files
committed
fix(reference): improve line attribution accuracy and add Claude Code model extraction
The reference implementation had two issues affecting trace accuracy: 1. Line Attribution: When processing edits with context lines (common in Claude Code's Edit tool), the entire new_string was attributed to AI, including unchanged surrounding lines. This produced inflated attribution ranges. 2. Model Identification: Claude Code does not include the model identifier in hook payloads. Traces were created with missing model_id, making it impossible to distinguish which model produced the code. Changes: - Add diffToFindChangedLines() to compute actual changed lines by comparing old_string and new_string, excluding context lines from attribution - Add extractModelFromTranscript() to parse Claude Code's JSONL transcript files and extract the model identifier from message entries - Add resolveModel() helper to transparently handle model resolution for both Cursor (direct payload) and Claude Code (transcript extraction) - Update PostToolUse, SessionStart, and SessionEnd handlers to use the new model resolution logic
1 parent e65936d commit 7962a28

2 files changed

Lines changed: 190 additions & 10 deletions

File tree

reference/trace-hook.ts

Lines changed: 35 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,23 @@
11
#!/usr/bin/env bun
22

3+
/**
4+
* Agent Trace Hook Handler
5+
*
6+
* This script processes hook events from AI coding tools (Cursor, Claude Code)
7+
* and generates trace records for attribution tracking. It reads JSON input
8+
* from stdin and dispatches to the appropriate handler based on hook_event_name.
9+
*
10+
* Supported tools:
11+
* - Cursor: afterFileEdit, afterTabFileEdit, afterShellExecution, sessionStart, sessionEnd
12+
* - Claude Code: PostToolUse, SessionStart, SessionEnd
13+
*/
14+
315
import {
416
createTrace,
517
appendTrace,
618
computeRangePositions,
719
tryReadFile,
8-
type ContributorType,
20+
extractModelFromTranscript,
921
type FileEdit,
1022
} from "./trace-store";
1123

@@ -32,6 +44,25 @@ interface HookInput {
3244
cwd?: string;
3345
}
3446

47+
/**
48+
* Resolves the model identifier from hook input.
49+
*
50+
* Different tools provide model information differently:
51+
* - Cursor: Sends model directly in the hook payload via `input.model`
52+
* - Claude Code: Does not include model in payload; must be extracted from transcript
53+
*
54+
* This function handles both cases transparently.
55+
*/
56+
function resolveModel(input: HookInput): string | undefined {
57+
if (input.model) {
58+
return input.model;
59+
}
60+
if (input.transcript_path) {
61+
return extractModelFromTranscript(input.transcript_path);
62+
}
63+
return undefined;
64+
}
65+
3566
const handlers: Record<string, (input: HookInput) => void> = {
3667
afterFileEdit: (input) => {
3768
const rangePositions = computeRangePositions(input.edits ?? [], tryReadFile(input.file_path!));
@@ -108,7 +139,7 @@ const handlers: Record<string, (input: HookInput) => void> = {
108139
: undefined;
109140

110141
appendTrace(createTrace("ai", file, {
111-
model: input.model,
142+
model: resolveModel(input),
112143
rangePositions,
113144
transcript: input.transcript_path,
114145
metadata: {
@@ -122,14 +153,14 @@ const handlers: Record<string, (input: HookInput) => void> = {
122153

123154
SessionStart: (input) => {
124155
appendTrace(createTrace("ai", ".sessions", {
125-
model: input.model,
156+
model: resolveModel(input),
126157
metadata: { event: "session_start", session_id: input.session_id, source: input.source },
127158
}));
128159
},
129160

130161
SessionEnd: (input) => {
131162
appendTrace(createTrace("ai", ".sessions", {
132-
model: input.model,
163+
model: resolveModel(input),
133164
metadata: { event: "session_end", session_id: input.session_id, reason: input.reason },
134165
}));
135166
},

reference/trace-store.ts

Lines changed: 155 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
import { execFileSync } from "child_process";
2-
import { existsSync, mkdirSync, appendFileSync, readFileSync } from "fs";
2+
import { existsSync, mkdirSync, appendFileSync, readFileSync, openSync, fstatSync, readSync, closeSync } from "fs";
33
import { join, relative } from "path";
44

55
export interface Range {
@@ -94,30 +94,179 @@ export function normalizeModelId(model?: string): string | undefined {
9494
return model;
9595
}
9696

97+
/**
98+
* Extracts the model identifier from a Claude Code transcript file.
99+
*
100+
* Claude Code stores conversation transcripts as JSONL files where each line
101+
* represents a message exchange. The model identifier is stored at `entry.message.model`.
102+
* This function reads only the tail of the file to efficiently get the most recent model,
103+
* which handles cases where the model may have changed during a session.
104+
*
105+
* @param transcriptPath - Absolute path to the Claude Code transcript JSONL file
106+
* @returns The model identifier (e.g., "claude-opus-4-5-20251101") or undefined if not found
107+
*
108+
* @example
109+
* ```typescript
110+
* const model = extractModelFromTranscript("/path/to/transcript.jsonl");
111+
* // Returns: "claude-opus-4-5-20251101"
112+
* ```
113+
*/
114+
export function extractModelFromTranscript(transcriptPath: string): string | undefined {
115+
try {
116+
const fd = openSync(transcriptPath, "r");
117+
const stats = fstatSync(fd);
118+
119+
// Read last 8KB - sufficient for recent JSONL entries
120+
const readSize = Math.min(stats.size, 8 * 1024);
121+
const buffer = Buffer.alloc(readSize);
122+
readSync(fd, buffer, 0, readSize, stats.size - readSize);
123+
closeSync(fd);
124+
125+
const content = buffer.toString("utf-8");
126+
const lines = content.split("\n");
127+
128+
// Iterate from end to get the most recent model
129+
for (let i = lines.length - 1; i >= 0; i--) {
130+
const line = lines[i].trim();
131+
if (!line) continue;
132+
133+
try {
134+
const entry = JSON.parse(line);
135+
if (entry.message?.model) {
136+
return entry.message.model;
137+
}
138+
} catch {
139+
// Skip malformed/partial JSON lines (first line may be truncated)
140+
continue;
141+
}
142+
}
143+
144+
return undefined;
145+
} catch {
146+
// File doesn't exist or isn't readable
147+
return undefined;
148+
}
149+
}
150+
97151
export interface RangePosition {
98152
start_line: number;
99153
end_line: number;
100154
}
101155

156+
/**
157+
* Computes which lines in `newStr` are actually new or modified compared to `oldStr`.
158+
*
159+
* This function performs a simple line-by-line diff to distinguish between:
160+
* - Context lines: Lines that exist in both old and new strings (not attributed)
161+
* - Changed lines: Lines that are new or modified (attributed to AI)
162+
*
163+
* This is necessary because some tools (like Claude Code's Edit tool) include
164+
* surrounding context lines in both `old_string` and `new_string`. Without this
165+
* diff, we would incorrectly attribute unchanged context lines to the AI.
166+
*
167+
* @param oldStr - The original string before the edit
168+
* @param newStr - The new string after the edit
169+
* @returns Array of 0-indexed line offsets within `newStr` that are new or modified
170+
*
171+
* @example
172+
* ```typescript
173+
* // old: "line1\nline2\nline3"
174+
* // new: "line1\nNEW LINE\nline3"
175+
* diffToFindChangedLines(old, new); // Returns [1] - only the middle line changed
176+
* ```
177+
*/
178+
function diffToFindChangedLines(oldStr: string, newStr: string): number[] {
179+
const oldLines = oldStr.split("\n");
180+
const newLines = newStr.split("\n");
181+
const changedOffsets: number[] = [];
182+
183+
let oldIdx = 0;
184+
185+
for (let newIdx = 0; newIdx < newLines.length; newIdx++) {
186+
if (oldIdx < oldLines.length && oldLines[oldIdx] === newLines[newIdx]) {
187+
// Matching line - this is context, not a change
188+
oldIdx++;
189+
} else {
190+
// Check if this line from newStr exists later in oldStr (handles deletions)
191+
let foundAhead = false;
192+
for (let lookAhead = oldIdx; lookAhead < oldLines.length; lookAhead++) {
193+
if (oldLines[lookAhead] === newLines[newIdx]) {
194+
oldIdx = lookAhead + 1;
195+
foundAhead = true;
196+
break;
197+
}
198+
}
199+
200+
if (!foundAhead) {
201+
// Line is genuinely new or modified - attribute to AI
202+
changedOffsets.push(newIdx);
203+
}
204+
}
205+
}
206+
207+
return changedOffsets;
208+
}
209+
102210
export function computeRangePositions(edits: FileEdit[], fileContent?: string): RangePosition[] {
103211
return edits
104212
.filter((e) => e.new_string)
105-
.map((edit) => {
213+
.flatMap((edit) => {
214+
// Case 1: Has explicit range from tool → use it
106215
if (edit.range) {
107-
return {
216+
return [{
108217
start_line: edit.range.start_line_number,
109218
end_line: edit.range.end_line_number,
110-
};
219+
}];
220+
}
221+
222+
// Case 2: Has both old_string and new_string → diff them to find actual changes
223+
if (edit.old_string && edit.new_string && fileContent) {
224+
const idx = fileContent.indexOf(edit.new_string);
225+
if (idx !== -1) {
226+
const startLine = fileContent.substring(0, idx).split("\n").length;
227+
const changedOffsets = diffToFindChangedLines(edit.old_string, edit.new_string);
228+
229+
if (changedOffsets.length === 0) {
230+
return [];
231+
}
232+
233+
// Convert offsets to line ranges, merging adjacent lines
234+
const ranges: RangePosition[] = [];
235+
let rangeStart = changedOffsets[0];
236+
let rangeEnd = changedOffsets[0];
237+
238+
for (let i = 1; i < changedOffsets.length; i++) {
239+
if (changedOffsets[i] === rangeEnd + 1) {
240+
rangeEnd = changedOffsets[i];
241+
} else {
242+
ranges.push({
243+
start_line: startLine + rangeStart,
244+
end_line: startLine + rangeEnd,
245+
});
246+
rangeStart = changedOffsets[i];
247+
rangeEnd = changedOffsets[i];
248+
}
249+
}
250+
251+
ranges.push({
252+
start_line: startLine + rangeStart,
253+
end_line: startLine + rangeEnd,
254+
});
255+
256+
return ranges;
257+
}
111258
}
259+
260+
// Case 3: Fallback - attribute entire new_string (original behavior)
112261
const lineCount = edit.new_string.split("\n").length;
113262
if (fileContent) {
114263
const idx = fileContent.indexOf(edit.new_string);
115264
if (idx !== -1) {
116265
const startLine = fileContent.substring(0, idx).split("\n").length;
117-
return { start_line: startLine, end_line: startLine + lineCount - 1 };
266+
return [{ start_line: startLine, end_line: startLine + lineCount - 1 }];
118267
}
119268
}
120-
return { start_line: 1, end_line: lineCount };
269+
return [{ start_line: 1, end_line: lineCount }];
121270
});
122271
}
123272

0 commit comments

Comments
 (0)