diff --git a/app/cli/internal/trace/claude/hooks.go b/app/cli/internal/trace/claude/hooks.go index ff25e943b..7c46ea86b 100644 --- a/app/cli/internal/trace/claude/hooks.go +++ b/app/cli/internal/trace/claude/hooks.go @@ -35,7 +35,12 @@ const ( eventSessionStart = "SessionStart" eventPreToolUse = "PreToolUse" eventPostToolUse = "PostToolUse" - eventSessionEnd = "SessionEnd" + // eventPostToolUseFailure fires instead of PostToolUse when a tool call + // fails. A failed shell command can still have changed files (e.g. a + // script that writes files and then runs a failing linter), so it runs + // the same handler. + eventPostToolUseFailure = "PostToolUseFailure" + eventSessionEnd = "SessionEnd" ) // fileWritingTools is the single source of truth for Claude tool names that modify files. @@ -61,6 +66,7 @@ var hookEvents = []hookEvent{ {eventSessionStart, "chainloop trace hook claude session-start", ""}, {eventPreToolUse, "chainloop trace hook claude pre-tool-use", hookToolMatcher}, {eventPostToolUse, "chainloop trace hook claude post-tool-use", hookToolMatcher}, + {eventPostToolUseFailure, "chainloop trace hook claude post-tool-use", hookToolMatcher}, {eventSessionEnd, "chainloop trace hook claude session-end", ""}, } @@ -194,6 +200,7 @@ func (p *Provider) ReadHookInput(r io.Reader) (*trace.HookInput, error) { if raw.ToolInput.FilePath != "" { input.FilePath = raw.ToolInput.FilePath } + input.ToolFailed = input.HookEventName == eventPostToolUseFailure return &input, nil } diff --git a/app/cli/internal/trace/claude/hooks_test.go b/app/cli/internal/trace/claude/hooks_test.go index c99b1072c..382645045 100644 --- a/app/cli/internal/trace/claude/hooks_test.go +++ b/app/cli/internal/trace/claude/hooks_test.go @@ -18,6 +18,7 @@ package claude import ( "bytes" "encoding/json" + "fmt" "os" "path/filepath" "testing" @@ -45,6 +46,9 @@ func TestInstallHooks(t *testing.T) { assertHookCommand(t, hooks, "SessionStart", "chainloop trace hook claude session-start") assertHookCommand(t, hooks, "PreToolUse", "chainloop trace hook claude pre-tool-use") assertHookCommand(t, hooks, "PostToolUse", "chainloop trace hook claude post-tool-use") + // A failed tool call fires PostToolUseFailure instead of PostToolUse, + // and a failed shell command can still have changed files. + assertHookCommand(t, hooks, "PostToolUseFailure", "chainloop trace hook claude post-tool-use") }) t.Run("installs PreToolUse and PostToolUse with matchers", func(t *testing.T) { @@ -64,6 +68,11 @@ func TestInstallHooks(t *testing.T) { postEntry := postEntries[0].(map[string]any) assert.Equal(t, "Edit|Write|MultiEdit|Bash", postEntry["matcher"]) + // PostToolUseFailure should have the same matcher + failureEntries := hooks["PostToolUseFailure"].([]any) + failureEntry := failureEntries[0].(map[string]any) + assert.Equal(t, "Edit|Write|MultiEdit|Bash", failureEntry["matcher"]) + // SessionStart should NOT have matcher startEntries := hooks["SessionStart"].([]any) startEntry := startEntries[0].(map[string]any) @@ -203,6 +212,7 @@ func TestUninstallHooks(t *testing.T) { assert.Contains(t, hooks, "PostToolUse") assert.NotContains(t, hooks, "SessionStart") assert.NotContains(t, hooks, "PreToolUse") + assert.NotContains(t, hooks, "PostToolUseFailure") }) t.Run("noop when file does not exist", func(t *testing.T) { @@ -229,6 +239,7 @@ func TestReadHookInput(t *testing.T) { "session_id": "abc-123", "hook_event_name": "PreToolUse", "tool_name": "Edit", + "tool_use_id": "toolu_01ABC", "tool_input": {"file_path": "/some/file.go", "old_string": "foo"} }`) input, err := provider.ReadHookInput(r) @@ -236,6 +247,7 @@ func TestReadHookInput(t *testing.T) { assert.Equal(t, "abc-123", input.SessionID) assert.Equal(t, "PreToolUse", input.HookEventName) assert.Equal(t, "Edit", input.ToolName) + assert.Equal(t, "toolu_01ABC", input.ToolUseID) assert.Equal(t, "/some/file.go", input.FilePath) }) @@ -248,6 +260,33 @@ func TestReadHookInput(t *testing.T) { assert.Empty(t, input.FilePath) }) + t.Run("flags failed tool calls", func(t *testing.T) { + testCases := []struct { + event string + wantFailed bool + }{ + {event: "PreToolUse", wantFailed: false}, + {event: "PostToolUse", wantFailed: false}, + {event: "PostToolUseFailure", wantFailed: true}, + } + + for _, tc := range testCases { + t.Run(tc.event, func(t *testing.T) { + r := bytes.NewBufferString(fmt.Sprintf(`{ + "session_id": "abc-123", + "hook_event_name": %q, + "tool_name": "Bash", + "tool_input": {"command": "make lint"}, + "error": "Exit code 1", + "is_interrupt": false + }`, tc.event)) + input, err := provider.ReadHookInput(r) + require.NoError(t, err) + assert.Equal(t, tc.wantFailed, input.ToolFailed) + }) + } + }) + t.Run("returns empty for empty session ID", func(t *testing.T) { r := bytes.NewBufferString(`{"session_id":""}`) input, err := provider.ReadHookInput(r) @@ -316,6 +355,7 @@ func TestInstallHooksForTraceRun(t *testing.T) { assert.Contains(t, hooks, "SessionStart") assert.Contains(t, hooks, "PreToolUse") assert.Contains(t, hooks, "PostToolUse") + assert.Contains(t, hooks, "PostToolUseFailure") assert.NotContains(t, hooks, "SessionEnd", "trace run must not install SessionEnd; trace run drives end-of-session itself") }) } diff --git a/app/cli/internal/trace/opencode/hooks.go b/app/cli/internal/trace/opencode/hooks.go index 69612b6de..da1c3a690 100644 --- a/app/cli/internal/trace/opencode/hooks.go +++ b/app/cli/internal/trace/opencode/hooks.go @@ -141,10 +141,13 @@ export const ChainloopTrace: Plugin = async ({ $, client }) => { }, "tool.execute.before": async (input, output) => { if (commandTools.includes(input.tool)) { + // callID pairs this hook with the tool.execute.after of the same + // call, so overlapping commands keep their own snapshots. await fire("pre-tool-use", { session_id: input.sessionID, hook_event_name: "tool.execute.before", tool_name: input.tool, + tool_use_id: input.callID, }) return } @@ -164,6 +167,7 @@ export const ChainloopTrace: Plugin = async ({ $, client }) => { session_id: input.sessionID, hook_event_name: "tool.execute.after", tool_name: input.tool, + tool_use_id: input.callID, }) return } @@ -267,6 +271,7 @@ func (p *Provider) ReadHookInput(r io.Reader) (*trace.HookInput, error) { HookEventName string `json:"hook_event_name"` ToolName string `json:"tool_name"` FilePath string `json:"file_path"` + ToolUseID string `json:"tool_use_id"` } if err := json.Unmarshal(data, &raw); err != nil { return nil, err @@ -277,6 +282,7 @@ func (p *Provider) ReadHookInput(r io.Reader) (*trace.HookInput, error) { HookEventName: raw.HookEventName, ToolName: raw.ToolName, FilePath: raw.FilePath, + ToolUseID: raw.ToolUseID, }, nil } diff --git a/app/cli/internal/trace/opencode/hooks_test.go b/app/cli/internal/trace/opencode/hooks_test.go index 175b1ab6b..33ee7dd99 100644 --- a/app/cli/internal/trace/opencode/hooks_test.go +++ b/app/cli/internal/trace/opencode/hooks_test.go @@ -139,6 +139,16 @@ func TestReadHookInputParsesValidInput(t *testing.T) { assert.Equal(t, "/some/file.go", input.FilePath) } +// The plugin sends opencode's callID as tool_use_id on shell hooks, so that +// the pre and post hooks of one command pair their snapshots. +func TestReadHookInputParsesToolUseID(t *testing.T) { + r := bytes.NewBufferString(`{"session_id":"ses_1","hook_event_name":"tool.execute.before","tool_name":"bash","tool_use_id":"call_01"}`) + p := New() + input, err := p.ReadHookInput(r) + require.NoError(t, err) + assert.Equal(t, "call_01", input.ToolUseID) +} + func TestReadHookInputApplyPatchSingleFile(t *testing.T) { // apply_patch fires one hook per file, so each invocation still carries // a single file_path — this is the shape the plugin emits after the fix. diff --git a/app/cli/internal/trace/opencode/testdata/plugin_full.ts b/app/cli/internal/trace/opencode/testdata/plugin_full.ts index 304e42c29..73c071ecc 100644 --- a/app/cli/internal/trace/opencode/testdata/plugin_full.ts +++ b/app/cli/internal/trace/opencode/testdata/plugin_full.ts @@ -85,10 +85,13 @@ export const ChainloopTrace: Plugin = async ({ $, client }) => { }, "tool.execute.before": async (input, output) => { if (commandTools.includes(input.tool)) { + // callID pairs this hook with the tool.execute.after of the same + // call, so overlapping commands keep their own snapshots. await fire("pre-tool-use", { session_id: input.sessionID, hook_event_name: "tool.execute.before", tool_name: input.tool, + tool_use_id: input.callID, }) return } @@ -108,6 +111,7 @@ export const ChainloopTrace: Plugin = async ({ $, client }) => { session_id: input.sessionID, hook_event_name: "tool.execute.after", tool_name: input.tool, + tool_use_id: input.callID, }) return } diff --git a/app/cli/internal/trace/opencode/testdata/plugin_tracerun.ts b/app/cli/internal/trace/opencode/testdata/plugin_tracerun.ts index ef5989db7..9626052be 100644 --- a/app/cli/internal/trace/opencode/testdata/plugin_tracerun.ts +++ b/app/cli/internal/trace/opencode/testdata/plugin_tracerun.ts @@ -81,10 +81,13 @@ export const ChainloopTrace: Plugin = async ({ $, client }) => { }, "tool.execute.before": async (input, output) => { if (commandTools.includes(input.tool)) { + // callID pairs this hook with the tool.execute.after of the same + // call, so overlapping commands keep their own snapshots. await fire("pre-tool-use", { session_id: input.sessionID, hook_event_name: "tool.execute.before", tool_name: input.tool, + tool_use_id: input.callID, }) return } @@ -104,6 +107,7 @@ export const ChainloopTrace: Plugin = async ({ $, client }) => { session_id: input.sessionID, hook_event_name: "tool.execute.after", tool_name: input.tool, + tool_use_id: input.callID, }) return } diff --git a/app/cli/internal/trace/provider.go b/app/cli/internal/trace/provider.go index 379144dd3..bb099d80f 100644 --- a/app/cli/internal/trace/provider.go +++ b/app/cli/internal/trace/provider.go @@ -179,6 +179,10 @@ type HookInput struct { // their parent's SessionID, so this is what tells concurrent agents of // one session apart. Empty for the main agent. AgentID string `json:"agent_id,omitempty"` + // ToolUseID is the agent's identifier for one tool call. Its pre and post + // hooks carry the same value, so it tells overlapping calls of one agent + // apart. Empty when the agent does not report it. + ToolUseID string `json:"tool_use_id,omitempty"` // AgentVersion is the agent runtime version reported in the hook payload // (e.g., Cursor's cursor_version). Captured at session-start so parsing // can set Agent.Version even when the transcript itself doesn't carry it. @@ -192,6 +196,11 @@ type HookInput struct { // only emit post-edit events (e.g., Cursor's afterFileEdit) populate it so // consumers can reconstruct the "before" content via reverse application. Edits []HookEdit `json:"-"` + // ToolFailed reports that the hook fires after a tool call that failed + // (Claude's PostToolUseFailure). The tool can still have changed files, + // so its changes are recorded. But the agent expects a different hook + // response for a failure, so nothing is written back to it. + ToolFailed bool `json:"-"` } // HookEdit represents a single old_string → new_string replacement applied to a file. diff --git a/app/cli/internal/trace/state/session.go b/app/cli/internal/trace/state/session.go index 76246ce1d..02bbd0dd3 100644 --- a/app/cli/internal/trace/state/session.go +++ b/app/cli/internal/trace/state/session.go @@ -44,6 +44,12 @@ type SessionRecord struct { // over Cwd to find the transcripts. Filled in by a later hook when the // first one did not carry it. Empty when the agent does not report it. TranscriptPath string `json:"transcript_path,omitempty"` + // Checkouts lists the roots of the other checkouts that this session + // edited with a file tool, sorted. Shell hooks run in the session's own + // checkout, and they also snapshot these checkouts, so that a command + // like `cd && …` is attributed there. Only set on the + // record in the session's own checkout. + Checkouts []string `json:"checkouts,omitempty"` // Active reports whether the session is ongoing at the time of record write. Active bool `json:"active"` // StartedAt is the RFC3339 timestamp of when tracking began for this session. diff --git a/app/cli/internal/trace/state/snapshot.go b/app/cli/internal/trace/state/snapshot.go index cde0a1da6..2be3411d6 100644 --- a/app/cli/internal/trace/state/snapshot.go +++ b/app/cli/internal/trace/state/snapshot.go @@ -52,26 +52,48 @@ func (s *Store) DeleteFileSnapshot(sessionID, filePath string) { _ = os.Remove(path) } +// ShellCallKey identifies the shell command that a pre-command signature +// belongs to, so the post-command hook of that command finds it. +type ShellCallKey struct { + SessionID string + // AgentID is empty for the main agent. Subagents share their parent's + // session ID, so this keeps their signatures apart. + AgentID string + // ToolUseID is the agent's identifier for the tool call. Empty when the + // agent does not report one. Then the agent has one slot, and its + // overlapping commands overwrite each other's signature. + ToolUseID string +} + // shellPreSignaturePath returns the path storing the pre-command working-tree -// signature for an agent of a session: -// /chainloop-trace/snapshots//shell-pre.json for the main agent, -// /chainloop-trace/snapshots//shell-pre-.json for a subagent. -func (s *Store) shellPreSignaturePath(sessionID, agentID string) string { +// signature of a shell command, under +// /chainloop-trace/snapshots//: +// - shell-pre-call-.json when the call has an ID; +// - shell-pre.json for the main agent otherwise; +// - shell-pre-.json for a subagent otherwise. +// +// A tool use ID is unique within a session, so it needs no agent qualifier. +func (s *Store) shellPreSignaturePath(key ShellCallKey) string { name := "shell-pre.json" - if agentID != "" { - name = "shell-pre-" + sanitizeID(agentID) + ".json" + switch { + case key.ToolUseID != "": + name = "shell-pre-call-" + sanitizeID(key.ToolUseID) + ".json" + case key.AgentID != "": + name = "shell-pre-" + sanitizeID(key.AgentID) + ".json" } - return filepath.Join(s.traceDirPath(), traceDirSnapshots, sanitizeID(sessionID), name) + return filepath.Join(s.traceDirPath(), traceDirSnapshots, sanitizeID(key.SessionID), name) } -// SaveShellPreSignature stores the working-tree signature captured before an -// agent-run shell command, so the post-command hook can diff against it. -// agentID is empty for the main agent. Subagents share their parent's session -// ID, so each agent gets its own slot; concurrent shell calls of one agent in -// one turn still overwrite it (see the parallel-shell limitation). -func (s *Store) SaveShellPreSignature(sessionID, agentID string, sig map[string]string) error { - path := s.shellPreSignaturePath(sessionID, agentID) +// WorktreeSignatures holds the working-tree signatures that a shell command +// is diffed against: checkout root → (repo-relative path → content hash). A +// command can change files in each checkout that its session edits. +type WorktreeSignatures map[string]map[string]string + +// SaveShellPreSignature stores the working-tree signatures captured before an +// agent-run shell command, so the post-command hook can diff against them. +func (s *Store) SaveShellPreSignature(key ShellCallKey, sig WorktreeSignatures) error { + path := s.shellPreSignaturePath(key) if err := os.MkdirAll(filepath.Dir(path), 0755); err != nil { return fmt.Errorf("create snapshot dir: %w", err) } @@ -84,15 +106,16 @@ func (s *Store) SaveShellPreSignature(sessionID, agentID string, sig map[string] return os.WriteFile(path, data, 0600) } -// LoadShellPreSignature loads the pre-command working-tree signature for an -// agent of a session. -func (s *Store) LoadShellPreSignature(sessionID, agentID string) (map[string]string, error) { - data, err := os.ReadFile(s.shellPreSignaturePath(sessionID, agentID)) +// LoadShellPreSignature loads the pre-command working-tree signatures of a +// shell command. A file in the earlier single-checkout format (written by an +// older CLI just before an upgrade) does not parse, and the caller skips it. +func (s *Store) LoadShellPreSignature(key ShellCallKey) (WorktreeSignatures, error) { + data, err := os.ReadFile(s.shellPreSignaturePath(key)) if err != nil { return nil, err } - var sig map[string]string + var sig WorktreeSignatures if err := json.Unmarshal(data, &sig); err != nil { return nil, fmt.Errorf("parse shell signature: %w", err) } @@ -101,6 +124,6 @@ func (s *Store) LoadShellPreSignature(sessionID, agentID string) (map[string]str } // DeleteShellPreSignature removes the pre-command signature once processed. -func (s *Store) DeleteShellPreSignature(sessionID, agentID string) { - _ = os.Remove(s.shellPreSignaturePath(sessionID, agentID)) +func (s *Store) DeleteShellPreSignature(key ShellCallKey) { + _ = os.Remove(s.shellPreSignaturePath(key)) } diff --git a/app/cli/internal/trace/state/snapshot_test.go b/app/cli/internal/trace/state/snapshot_test.go index 471be6a03..ef262d656 100644 --- a/app/cli/internal/trace/state/snapshot_test.go +++ b/app/cli/internal/trace/state/snapshot_test.go @@ -16,67 +16,99 @@ package state import ( + "fmt" + "os" + "path/filepath" "testing" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) +const shellSessionID = "sess-123" + func TestShellPreSignatureRoundTrip(t *testing.T) { cases := []struct { - name string - agentID string + name string + key ShellCallKey }{ - {name: "main session", agentID: ""}, - {name: "subagent", agentID: "afd65659e2015d48d"}, + {name: "main session", key: ShellCallKey{SessionID: shellSessionID}}, + {name: "subagent", key: ShellCallKey{SessionID: shellSessionID, AgentID: "afd65659e2015d48d"}}, + {name: "tool call", key: ShellCallKey{SessionID: shellSessionID, ToolUseID: "toolu_01ABC"}}, + {name: "subagent tool call", key: ShellCallKey{SessionID: shellSessionID, AgentID: "afd65659e2015d48d", ToolUseID: "toolu_01ABC"}}, } for _, tc := range cases { t.Run(tc.name, func(t *testing.T) { store := NewGitStore(t.TempDir()) - sessionID := "sess-123" - sig := map[string]string{ - "a.go": "hash-a", - "sub/b.json": "hash-b", + sig := WorktreeSignatures{ + "/repo/home": { + fileA: "hash-a", + "sub/b.json": "hash-b", + }, + "/repo/other": { + "c.ts": "hash-c", + }, } - require.NoError(t, store.SaveShellPreSignature(sessionID, tc.agentID, sig)) + require.NoError(t, store.SaveShellPreSignature(tc.key, sig)) - loaded, err := store.LoadShellPreSignature(sessionID, tc.agentID) + loaded, err := store.LoadShellPreSignature(tc.key) require.NoError(t, err) assert.Equal(t, sig, loaded) - store.DeleteShellPreSignature(sessionID, tc.agentID) + store.DeleteShellPreSignature(tc.key) - _, err = store.LoadShellPreSignature(sessionID, tc.agentID) + _, err = store.LoadShellPreSignature(tc.key) assert.Error(t, err, "signature should be gone after delete") }) } } -// A subagent shares its parent's session ID. Their shell commands can overlap, -// so each agent needs its own slot or one deletes the other's signature. -func TestShellPreSignaturePerAgent(t *testing.T) { +// Shell commands can overlap: a subagent shares its parent's session ID, and +// one agent can run several commands at once. Each command needs its own slot, +// or the first post hook to finish deletes a signature another command still +// needs. +func TestShellPreSignatureSlots(t *testing.T) { + keys := []ShellCallKey{ + {SessionID: shellSessionID}, + {SessionID: shellSessionID, AgentID: "agent-1"}, + {SessionID: shellSessionID, AgentID: "agent-2"}, + {SessionID: shellSessionID, ToolUseID: "toolu_1"}, + {SessionID: shellSessionID, ToolUseID: "toolu_2"}, + {SessionID: shellSessionID, AgentID: "agent-1", ToolUseID: "toolu_3"}, + } + store := NewGitStore(t.TempDir()) - const sessionID = "sess-123" - parent := map[string]string{"a.go": "parent"} - sub1 := map[string]string{"a.go": "sub1"} - sub2 := map[string]string{"a.go": "sub2"} + for i, key := range keys { + require.NoError(t, store.SaveShellPreSignature(key, WorktreeSignatures{"/repo": {fileA: fmt.Sprint(i)}})) + } - require.NoError(t, store.SaveShellPreSignature(sessionID, "", parent)) - require.NoError(t, store.SaveShellPreSignature(sessionID, "agent-1", sub1)) - require.NoError(t, store.SaveShellPreSignature(sessionID, "agent-2", sub2)) + deleted := keys[1] + store.DeleteShellPreSignature(deleted) - store.DeleteShellPreSignature(sessionID, "agent-1") + for i, key := range keys { + got, err := store.LoadShellPreSignature(key) + if key == deleted { + assert.Error(t, err, "deleted slot %+v", key) + continue + } - got, err := store.LoadShellPreSignature(sessionID, "") - require.NoError(t, err) - assert.Equal(t, parent, got) + require.NoError(t, err, "slot %+v", key) + assert.Equal(t, WorktreeSignatures{"/repo": {fileA: fmt.Sprint(i)}}, got, "slot %+v keeps its own signature", key) + } +} + +// A CLI upgrade can land between the pre and post hooks of one command. The +// earlier single-checkout format must then be rejected, not misread. +func TestShellPreSignatureRejectsSingleCheckoutFormat(t *testing.T) { + store := NewGitStore(t.TempDir()) + key := ShellCallKey{SessionID: shellSessionID} - got, err = store.LoadShellPreSignature(sessionID, "agent-2") - require.NoError(t, err) - assert.Equal(t, sub2, got) + path := store.shellPreSignaturePath(key) + require.NoError(t, os.MkdirAll(filepath.Dir(path), 0755)) + require.NoError(t, os.WriteFile(path, []byte(`{"a.go":"hash-a"}`), 0600)) - _, err = store.LoadShellPreSignature(sessionID, "agent-1") + _, err := store.LoadShellPreSignature(key) assert.Error(t, err) } diff --git a/app/cli/pkg/action/trace_agent_hook.go b/app/cli/pkg/action/trace_agent_hook.go index b030db01a..afac2b2cf 100644 --- a/app/cli/pkg/action/trace_agent_hook.go +++ b/app/cli/pkg/action/trace_agent_hook.go @@ -20,6 +20,7 @@ import ( "errors" "os" "path/filepath" + "slices" "strings" "github.com/chainloop-dev/chainloop/app/cli/internal/repositoryconfig" @@ -243,7 +244,7 @@ func HandleAgentPreToolUse(provider trace.Provider, log zerolog.Logger) error { case provider.IsCommandTool(input.ToolName): // Shell command: snapshot the whole worktree so the post hook can diff // it and attribute the command's file changes to the AI. - captureWorktreeSnapshot(store, repoRoot, input.SessionID, input.AgentID, log) + captureWorktreeSnapshot(store, repoRoot, shellCallKey(input), log) case provider.IsFileWritingTool(input.ToolName): if input.FilePath == "" { log.Debug().Str("tool", input.ToolName).Msg("pre-tool-use: file-writing tool produced no file path, skipping") @@ -252,6 +253,7 @@ func HandleAgentPreToolUse(provider trace.Provider, log zerolog.Logger) error { if err := provider.CaptureFileSnapshot(store, input); err != nil { log.Debug().Err(err).Str("file", input.FilePath).Msg("pre-tool-use: capture snapshot failed") } + registerSessionCheckout(input.SessionID, repoRoot, log) default: log.Debug().Str("tool", input.ToolName).Msg("pre-tool-use: not a tracked tool, skipping") } @@ -259,22 +261,109 @@ func HandleAgentPreToolUse(provider trace.Provider, log zerolog.Logger) error { return nil } -// captureWorktreeSnapshot records the working-tree signature before a shell -// command runs, so HandleAgentPostToolUse can diff it and attribute the files -// the command changed. agentID is empty for the main agent. Best-effort: -// failures are logged and never block the agent. -func captureWorktreeSnapshot(store *state.Store, repoRoot, sessionID, agentID string, log zerolog.Logger) { - sig, err := tracegit.NewGoGitClient().SnapshotWorktree(repoRoot) - if err != nil { - log.Debug().Err(err).Msg("pre-command: worktree snapshot failed") +// shellCallKey returns the key that pairs the pre and post hooks of one shell +// command. +func shellCallKey(input *trace.HookInput) state.ShellCallKey { + return state.ShellCallKey{SessionID: input.SessionID, AgentID: input.AgentID, ToolUseID: input.ToolUseID} +} + +// captureWorktreeSnapshot records the working-tree signatures before a shell +// command runs, so HandleAgentPostToolUse can diff them and attribute the +// files the command changed. It snapshots the session's own checkout and +// every other checkout that the session edited with a file tool: a command +// like `cd && …` changes files there, and shell hooks only +// run in the session's own checkout. Best-effort: failures are logged and +// never block the agent. +func captureWorktreeSnapshot(store *state.Store, repoRoot string, key state.ShellCallKey, log zerolog.Logger) { + client := tracegit.NewGoGitClient() + sigs := make(state.WorktreeSignatures) + for _, root := range shellSnapshotRoots(store, repoRoot, key.SessionID) { + sig, err := client.SnapshotWorktree(root) + if err != nil { + log.Debug().Err(err).Str("root", root).Msg("pre-command: worktree snapshot failed") + continue + } + sigs[root] = sig + } + + if len(sigs) == 0 { return } - if err := store.SaveShellPreSignature(sessionID, agentID, sig); err != nil { + if err := store.SaveShellPreSignature(key, sigs); err != nil { log.Debug().Err(err).Msg("pre-command: save worktree signature failed") } } +// shellSnapshotRoots returns the checkouts that a shell command of the session +// is diffed in: the session's own checkout, then the other checkouts on its +// session record that still exist. +func shellSnapshotRoots(store *state.Store, repoRoot, sessionID string) []string { + roots := []string{repoRoot} + + rec, err := store.LoadSessionRecord(sessionID) + if err != nil || rec == nil { + return roots + } + + for _, root := range rec.Checkouts { + if sameDir(root, repoRoot) { + continue + } + if _, err := os.Stat(root); err != nil { + continue + } + roots = append(roots, root) + } + + return roots +} + +// registerSessionCheckout adds fileRoot to the checkout list on the session +// record of the session's own checkout, when a file tool edits a file in a +// different checkout. Shell hooks then also snapshot that checkout. +// +// The record is never created here: it is missing only when the session's +// directory is not a repository (shell hooks record nothing then) or the hooks +// were installed partway through the session, and creating it would also copy +// the transcripts into that checkout. Best-effort: failures are logged and +// never block the agent. +func registerSessionCheckout(sessionID, fileRoot string, log zerolog.Logger) { + homeStore, homeRoot, err := state.Locate() + if err != nil || sameDir(homeRoot, fileRoot) { + return + } + + rec, err := homeStore.LoadSessionRecord(sessionID) + if err != nil || rec == nil { + log.Debug().Err(err).Str("root", fileRoot).Msg("no session record in the session's checkout; shell commands will not snapshot this checkout") + return + } + + if slices.Contains(rec.Checkouts, fileRoot) { + return + } + + rec.Checkouts = append(rec.Checkouts, fileRoot) + slices.Sort(rec.Checkouts) + if err := homeStore.SaveSessionRecord(rec); err != nil { + log.Debug().Err(err).Str("root", fileRoot).Msg("register session checkout failed") + } +} + +// sameDir reports whether a and b name the same directory, ignoring symlinks +// (e.g. /var and /private/var on macOS). +func sameDir(a, b string) bool { + if filepath.Clean(a) == filepath.Clean(b) { + return true + } + + ra, errA := filepath.EvalSymlinks(a) + rb, errB := filepath.EvalSymlinks(b) + + return errA == nil && errB == nil && ra == rb +} + // ensureSessionTracked auto-installs git hooks (unconditionally, on every // invocation — itself idempotent via hooks.IsInstalled), then creates a // session record if one doesn't already exist and copies the session data. @@ -476,17 +565,23 @@ func HandleAgentPostToolUse(provider trace.Provider, log zerolog.Logger) error { if isCommand { // Shell command: diff the before/after worktree snapshots and attribute // every file the command changed to the AI. - recordCommandLineRanges(store, repoRoot, sessionID, input.AgentID, log) + recordCommandLineRanges(provider, input, store, repoRoot, log) // The command may have been a `git push`, whose pre-push hook attested // a session and left its link behind. Show it now: the pre-push output // went to this tool call's captured stderr, which the user does not - // necessarily read. - notifyPendingSessionLinks(provider, store, log) + // necessarily read. A failed call expects a different hook response, + // so its links stay on disk for the next command that succeeds. + if !input.ToolFailed { + notifyPendingSessionLinks(provider, store, log) + } return nil } + // Cursor has no pre-tool-use, so the post hook registers the checkout too. + registerSessionCheckout(sessionID, repoRoot, log) + after, err := os.ReadFile(input.FilePath) if err != nil { if !errors.Is(err, os.ErrNotExist) { @@ -524,24 +619,47 @@ func HandleAgentPostToolUse(provider trace.Provider, log zerolog.Logger) error { return nil } -// recordCommandLineRanges diffs the worktree signature captured before a shell -// command against the current worktree, and records every created/modified file -// (whole-file range) and every deleted file as AI-attributed. Enrich later caps -// the AI line count to each file's committed diff totals, so whole-file ranges -// yield correct counts. agentID is empty for the main agent. Best-effort: -// never blocks the agent. -func recordCommandLineRanges(store *state.Store, repoRoot, sessionID, agentID string, log zerolog.Logger) { - before, err := store.LoadShellPreSignature(sessionID, agentID) +// recordCommandLineRanges diffs each worktree signature captured before a +// shell command against its current worktree, and records the changes in the +// ledger of the checkout that holds them. store and repoRoot are the +// session's own checkout, which holds the signatures. Best-effort: never +// blocks the agent. +func recordCommandLineRanges(provider trace.Provider, input *trace.HookInput, store *state.Store, repoRoot string, log zerolog.Logger) { + key := shellCallKey(input) + before, err := store.LoadShellPreSignature(key) if err != nil { - // No pre-command snapshot (missed pre hook, parallel overwrite) — skip. + // No pre-command snapshot (missed pre hook, or an overlapping call + // of an agent without tool call IDs overwrote it) — skip. log.Debug().Err(err).Msg("post-command: no pre-command worktree signature") return } - defer store.DeleteShellPreSignature(sessionID, agentID) + defer store.DeleteShellPreSignature(key) + + for root, rootBefore := range before { + rootStore := store + if !sameDir(root, repoRoot) { + s, r, err := state.LocateFrom(root) + if err != nil { + log.Debug().Err(err).Str("root", root).Msg("post-command: no trace state for checkout") + continue + } + ensureSessionTracked(provider, s, r, input, log) + rootStore, root = s, r + } + + recordRootChanges(rootStore, root, rootBefore, key.SessionID, log) + } +} +// recordRootChanges diffs the signature of one checkout captured before a +// shell command against its current worktree, and records every +// created/modified file (whole-file range) and every deleted file as +// AI-attributed. Enrich later caps the AI line count to each file's committed +// diff totals, so whole-file ranges yield correct counts. +func recordRootChanges(store *state.Store, repoRoot string, before map[string]string, sessionID string, log zerolog.Logger) { after, err := tracegit.NewGoGitClient().SnapshotWorktree(repoRoot) if err != nil { - log.Debug().Err(err).Msg("post-command: worktree snapshot failed") + log.Debug().Err(err).Str("root", repoRoot).Msg("post-command: worktree snapshot failed") return } @@ -568,7 +686,7 @@ func recordCommandLineRanges(store *state.Store, repoRoot, sessionID, agentID st } } - log.Debug().Int("changed", len(changed)).Int("deleted", len(deleted)).Str("session_id", sessionID).Msg("command line ranges recorded") + log.Debug().Int("changed", len(changed)).Int("deleted", len(deleted)).Str("session_id", sessionID).Str("root", repoRoot).Msg("command line ranges recorded") } // autoInstallGitHooks installs git hooks if they're not already present diff --git a/app/cli/pkg/action/trace_agent_hook_test.go b/app/cli/pkg/action/trace_agent_hook_test.go index 90730e208..b8de2a207 100644 --- a/app/cli/pkg/action/trace_agent_hook_test.go +++ b/app/cli/pkg/action/trace_agent_hook_test.go @@ -560,7 +560,7 @@ func TestHandleAgentCommandTool_AttributesShellFileChanges(t *testing.T) { assert.Empty(t, attr.Files["marker"]) // The pre-command signature is cleaned up afterwards. - _, err := store.LoadShellPreSignature("ses-cmd", "") + _, err := store.LoadShellPreSignature(state.ShellCallKey{SessionID: "ses-cmd"}) assert.Error(t, err) } @@ -635,7 +635,7 @@ func TestHandleAgentClaudeCodeSession(t *testing.T) { assert.Equal(t, 4, genRanges[0].End) // whole 4-line generated file // The pre-command signature is cleaned up after the Bash post hook. - _, err := store.LoadShellPreSignature(sid, "") + _, err := store.LoadShellPreSignature(state.ShellCallKey{SessionID: sid}) assert.Error(t, err) } @@ -674,11 +674,201 @@ func TestHandleAgentCommandTool_ConcurrentSubagent(t *testing.T) { assert.Contains(t, attr.Files, "parent.txt", "the parent's command must keep its own pre-command signature") for _, id := range []string{"", agentID} { - _, err := store.LoadShellPreSignature(sid, id) + _, err := store.LoadShellPreSignature(state.ShellCallKey{SessionID: sid, AgentID: id}) assert.Error(t, err, "signature for agent %q is cleaned up", id) } } +// One agent can run shell commands that overlap (for example, a background +// command and a foreground one). Each call carries its own tool_use_id, so +// each must keep its own pre-command signature: the second pre hook must not +// overwrite the first, and the first post hook must not delete the second's. +func TestHandleAgentCommandTool_OverlappingCallsOfOneAgent(t *testing.T) { + root := chdirToResolvedGitRepo(t) + store := state.NewGitStore(filepath.Join(root, ".git")) + require.NoError(t, store.InitTraceDir()) + + p := claude.New() + const sid = "e0e0e0c2-1a2b-4c3d-8e9f-0a1b2c3d4e5f" + + bash := func(event, toolUseID string, handler func(trace.Provider, zerolog.Logger) error) { + t.Helper() + withStdin(t, fmt.Sprintf(`{"session_id":%q,"cwd":%q,"hook_event_name":%q,"tool_name":"Bash","tool_use_id":%q,"tool_input":{"command":"gen"}}`, + sid, root, event, toolUseID)) + require.NoError(t, handler(p, zerolog.Nop())) + } + + bash("PreToolUse", "toolu_first", HandleAgentPreToolUse) + require.NoError(t, os.WriteFile(filepath.Join(root, "first.txt"), []byte("first\n"), 0600)) + + // The second command starts after the first one wrote its file. With one + // slot per agent, its snapshot replaces the first one and already + // contains first.txt. + bash("PreToolUse", "toolu_second", HandleAgentPreToolUse) + require.NoError(t, os.WriteFile(filepath.Join(root, "second.txt"), []byte("second\n"), 0600)) + + bash("PostToolUse", "toolu_first", HandleAgentPostToolUse) + bash("PostToolUse", "toolu_second", HandleAgentPostToolUse) + + attr := store.LoadAILineAttribution(sid) + assert.Contains(t, attr.Files, "first.txt", "the first command keeps its own pre-command signature") + assert.Contains(t, attr.Files, "second.txt", "the second command keeps its own pre-command signature") + + for _, id := range []string{"toolu_first", "toolu_second"} { + _, err := store.LoadShellPreSignature(state.ShellCallKey{SessionID: sid, ToolUseID: id}) + assert.Error(t, err, "signature for call %q is cleaned up", id) + } +} + +// A session can run in one checkout and edit another one: it writes files +// there with a file tool, and it changes more files there with shell commands +// like `cd && …`. Shell hooks run in the session's +// directory, so they must also snapshot every checkout that the session +// edited with a file tool, and record each change in the ledger of the +// checkout that holds the file. +func TestHandleAgentCommandTool_OtherCheckout(t *testing.T) { + home := chdirToResolvedGitRepo(t) + other := resolvedTempGitRepo(t) + untouched := resolvedTempGitRepo(t) + + homeStore := state.NewGitStore(filepath.Join(home, ".git")) + otherStore := state.NewGitStore(filepath.Join(other, ".git")) + untouchedStore := state.NewGitStore(filepath.Join(untouched, ".git")) + + p := claude.New() + const sid = "f0f0e0c2-1a2b-4c3d-8e9f-0a1b2c3d4e5f" + + hook := func(event, tool, toolUseID, toolInput string, handler func(trace.Provider, zerolog.Logger) error) { + t.Helper() + withStdin(t, fmt.Sprintf(`{"session_id":%q,"cwd":%q,"hook_event_name":%q,"tool_name":%q,"tool_use_id":%q,"tool_input":%s}`, + sid, home, event, tool, toolUseID, toolInput)) + require.NoError(t, handler(p, zerolog.Nop())) + } + + withStdin(t, fmt.Sprintf(`{"session_id":%q,"cwd":%q,"hook_event_name":"SessionStart","source":"startup"}`, sid, home)) + require.NoError(t, HandleAgentSessionStart(p, zerolog.Nop())) + + // 1. The agent writes a file in the other checkout, twice. + modal := filepath.Join(other, "modal.tsx") + writeInput := fmt.Sprintf(`{"file_path":%q,"content":"export {}\n"}`, modal) + for _, id := range []string{"toolu_w1", "toolu_w2"} { + hook("PreToolUse", "Write", id, writeInput, HandleAgentPreToolUse) + require.NoError(t, os.WriteFile(modal, []byte("export {}\n"), 0600)) + hook("PostToolUse", "Write", id, writeInput, HandleAgentPostToolUse) + } + + rec, err := homeStore.LoadSessionRecord(sid) + require.NoError(t, err) + require.NotNil(t, rec) + assert.Equal(t, []string{other}, rec.Checkouts, "the other checkout is registered once on the session record") + + // 2. A shell command changes files in the other checkout and in a + // checkout that no file tool touched. + bashInput := fmt.Sprintf(`{"command":"cd %s && python3 gen.py"}`, other) + hook("PreToolUse", "Bash", "toolu_b1", bashInput, HandleAgentPreToolUse) + require.NoError(t, os.WriteFile(filepath.Join(other, "init.txt"), []byte("changed\n"), 0600)) + require.NoError(t, os.MkdirAll(filepath.Join(other, "sheet"), 0755)) + require.NoError(t, os.WriteFile(filepath.Join(other, "sheet", "list.tsx"), []byte("a\nb\n"), 0600)) + require.NoError(t, os.WriteFile(filepath.Join(untouched, "w.go"), []byte("package w\n"), 0600)) + hook("PostToolUse", "Bash", "toolu_b1", bashInput, HandleAgentPostToolUse) + + otherAttr := otherStore.LoadAILineAttribution(sid) + assert.Contains(t, otherAttr.Files, "modal.tsx", "the file tool edit is recorded") + assert.Contains(t, otherAttr.Files, "init.txt", "a shell edit in the other checkout is recorded there") + require.Contains(t, otherAttr.Files, "sheet/list.tsx", "paths are relative to the checkout that holds the file") + assert.Equal(t, 2, otherAttr.Files["sheet/list.tsx"][0].End) + + homeAttr := homeStore.LoadAILineAttribution(sid) + assert.Empty(t, homeAttr.Files, "the session's own checkout did not change") + + // Accepted gap: without a file tool edit there, a checkout is not + // snapshotted. + assert.Empty(t, untouchedStore.LoadAILineAttribution(sid).Files) + assert.False(t, untouchedStore.SessionRecordExists(sid)) + + _, err = homeStore.LoadShellPreSignature(state.ShellCallKey{SessionID: sid, ToolUseID: "toolu_b1"}) + assert.Error(t, err, "the pre-command signature is cleaned up") +} + +// The checkout list lives on the session record of the session's own +// checkout. When that record is missing (the hooks were installed partway +// through the session), registration must not create one: creating it would +// also copy the transcripts into that checkout. +func TestHandleAgentFileTool_OtherCheckoutWithoutHomeRecord(t *testing.T) { + home := chdirToResolvedGitRepo(t) + other := resolvedTempGitRepo(t) + homeStore := state.NewGitStore(filepath.Join(home, ".git")) + + p := claude.New() + const sid = "a0a0e0c2-1a2b-4c3d-8e9f-0a1b2c3d4e5f" + + file := filepath.Join(other, "x.go") + writeInput := fmt.Sprintf(`"tool_name":"Write","tool_input":{"file_path":%q,"content":"package x\n"}`, file) + withStdin(t, fmt.Sprintf(`{"session_id":%q,"cwd":%q,"hook_event_name":"PreToolUse",%s}`, sid, home, writeInput)) + require.NoError(t, HandleAgentPreToolUse(p, zerolog.Nop())) + require.NoError(t, os.WriteFile(file, []byte("package x\n"), 0600)) + withStdin(t, fmt.Sprintf(`{"session_id":%q,"cwd":%q,"hook_event_name":"PostToolUse",%s}`, sid, home, writeInput)) + require.NoError(t, HandleAgentPostToolUse(p, zerolog.Nop())) + + assert.False(t, homeStore.SessionRecordExists(sid), "registration does not create the home record") + assert.Contains(t, state.NewGitStore(filepath.Join(other, ".git")).LoadAILineAttribution(sid).Files, "x.go") +} + +// resolvedTempGitRepo creates a git repo and returns its symlink-resolved +// root, without changing the working directory. +func resolvedTempGitRepo(t *testing.T) string { + t.Helper() + + resolved, err := filepath.EvalSymlinks(initTempGitRepo(t)) + require.NoError(t, err) + + return resolved +} + +// A shell command that fails still changes the files it wrote before it +// failed. Claude Code reports such a call through PostToolUseFailure, not +// PostToolUse, so the handler must attribute the command's changes from that +// payload too. +func TestHandleAgentCommandTool_FailedCommand(t *testing.T) { + root := chdirToResolvedGitRepo(t) + store := state.NewGitStore(filepath.Join(root, ".git")) + require.NoError(t, store.InitTraceDir()) + + p := claude.New() + const sid = "d1e2f3a4-1a2b-4c3d-8e9f-0a1b2c3d4e5f" + bashInput := `"tool_name":"Bash","tool_input":{"command":"python3 gen.py && npx eslint ."}` + + withStdin(t, fmt.Sprintf(`{"session_id":%q,"cwd":%q,"hook_event_name":"PreToolUse",%s}`, sid, root, bashInput)) + require.NoError(t, HandleAgentPreToolUse(p, zerolog.Nop())) + + // The command writes three files, then its linter step fails. + for _, name := range []string{"colors.ts", "check.tsx", "card.tsx"} { + require.NoError(t, os.WriteFile(filepath.Join(root, name), []byte("export {}\n"), 0600)) + } + + // A link left by an earlier push must wait for a successful command: the + // link announcement is a PostToolUse response, which does not match the + // failure event. + const link = "https://app.chainloop.dev/u/org/sessions/ses_1" + require.NoError(t, store.SavePendingLinks([]string{link})) + + withStdin(t, fmt.Sprintf(`{"session_id":%q,"cwd":%q,"hook_event_name":"PostToolUseFailure",%s,"error":"Exit code 1","is_interrupt":false}`, sid, root, bashInput)) + out := captureStdout(t, func() { + require.NoError(t, HandleAgentPostToolUse(p, zerolog.Nop())) + }) + + attr := store.LoadAILineAttribution(sid) + for _, name := range []string{"colors.ts", "check.tsx", "card.tsx"} { + assert.Contains(t, attr.Files, name, "a file written by a failed command is AI-made") + } + + _, err := store.LoadShellPreSignature(state.ShellCallKey{SessionID: sid}) + assert.Error(t, err, "the pre-command signature is cleaned up") + + assert.Empty(t, out, "a failed tool call gets no hook response") + assert.Equal(t, []string{link}, store.PendingLinks(), "the link stays for the next successful command") +} + // chdirToResolvedGitRepo creates a git repo, chdirs into its symlink-resolved // path, and returns that canonical root. Using the resolved path mirrors the // canonical absolute paths Claude Code passes in hook payloads and keeps