Skip to content

Commit f911ab2

Browse files
lancekrogerslancekrogers
committed
[OBEY-CAMPAIGN-78887e36] Fix plugin hook execution and budget enforcement
1 parent bc3f06f commit f911ab2

7 files changed

Lines changed: 360 additions & 142 deletions

File tree

‎README.md‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -221,8 +221,14 @@ pm.Register(filter, nil)
221221
audit := claude.NewAuditPlugin(1000) // Keep last 1000 records
222222
pm.Register(audit, nil)
223223

224+
ctx := context.Background()
225+
if err := pm.Initialize(ctx); err != nil {
226+
log.Fatal(err)
227+
}
228+
defer pm.Shutdown(ctx)
229+
224230
// Use with client
225-
result, err := cc.RunPrompt("Do something", &claude.RunOptions{
231+
result, err := cc.RunPromptCtx(ctx, "Do something", &claude.RunOptions{
226232
PluginManager: pm,
227233
})
228234

‎pkg/claude/claude.go‎

Lines changed: 16 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -57,6 +57,10 @@ func (c *ClaudeClient) RunPromptCtx(ctx context.Context, prompt string, opts *Ru
5757
opts = c.DefaultOptions
5858
}
5959

60+
if opts.PluginManager != nil && opts.Format == JSONOutput {
61+
return c.runPromptWithStructuredHooks(ctx, prompt, nil, opts)
62+
}
63+
6064
// Preprocess and validate options
6165
if err := PreprocessOptions(opts); err != nil {
6266
return nil, err
@@ -71,11 +75,9 @@ func (c *ClaudeClient) RunPromptCtx(ctx context.Context, prompt string, opts *Ru
7175

7276
args := BuildArgs(prompt, opts)
7377

74-
cleanupPlugins, err := preparePluginManager(ctx, opts.PluginManager)
75-
if err != nil {
78+
if err := ensurePluginManagerInitialized(ctx, opts.PluginManager); err != nil {
7679
return nil, err
7780
}
78-
defer cleanupPlugins()
7981

8082
cmd := execCommand(ctx, c.BinPath, args...)
8183
if opts.WorkingDirectory != "" {
@@ -85,7 +87,7 @@ func (c *ClaudeClient) RunPromptCtx(ctx context.Context, prompt string, opts *Ru
8587
cmd.Stdout = &stdout
8688
cmd.Stderr = &stderr
8789

88-
err = cmd.Run()
90+
err := cmd.Run()
8991
if err != nil {
9092
// Enhanced error parsing
9193
var exitCode int
@@ -106,7 +108,7 @@ func (c *ClaudeClient) RunPromptCtx(ctx context.Context, prompt string, opts *Ru
106108
return nil, err
107109
}
108110
if err := applyExecutionHooks(ctx, opts, messages, result); err != nil {
109-
return nil, err
111+
return result, err
110112
}
111113
return result, nil
112114
}
@@ -117,7 +119,7 @@ func (c *ClaudeClient) RunPromptCtx(ctx context.Context, prompt string, opts *Ru
117119
IsError: false,
118120
}
119121
if err := applyCompletionHooks(ctx, opts, result); err != nil {
120-
return nil, err
122+
return result, err
121123
}
122124
return result, nil
123125
}
@@ -139,6 +141,10 @@ func (c *ClaudeClient) RunFromStdinCtx(ctx context.Context, stdin io.Reader, pro
139141
opts = c.DefaultOptions
140142
}
141143

144+
if opts.PluginManager != nil && opts.Format == JSONOutput {
145+
return c.runPromptWithStructuredHooks(ctx, prompt, stdin, opts)
146+
}
147+
142148
// Preprocess and validate options
143149
if err := PreprocessOptions(opts); err != nil {
144150
return nil, err
@@ -153,11 +159,9 @@ func (c *ClaudeClient) RunFromStdinCtx(ctx context.Context, stdin io.Reader, pro
153159

154160
args := BuildArgs(prompt, opts)
155161

156-
cleanupPlugins, err := preparePluginManager(ctx, opts.PluginManager)
157-
if err != nil {
162+
if err := ensurePluginManagerInitialized(ctx, opts.PluginManager); err != nil {
158163
return nil, err
159164
}
160-
defer cleanupPlugins()
161165

162166
cmd := execCommand(ctx, c.BinPath, args...)
163167
if opts.WorkingDirectory != "" {
@@ -168,7 +172,7 @@ func (c *ClaudeClient) RunFromStdinCtx(ctx context.Context, stdin io.Reader, pro
168172
cmd.Stdout = &stdout
169173
cmd.Stderr = &stderr
170174

171-
err = cmd.Run()
175+
err := cmd.Run()
172176
if err != nil {
173177
// Enhanced error parsing
174178
var exitCode int
@@ -189,7 +193,7 @@ func (c *ClaudeClient) RunFromStdinCtx(ctx context.Context, stdin io.Reader, pro
189193
return nil, err
190194
}
191195
if err := applyExecutionHooks(ctx, opts, messages, result); err != nil {
192-
return nil, err
196+
return result, err
193197
}
194198
return result, nil
195199
}
@@ -200,7 +204,7 @@ func (c *ClaudeClient) RunFromStdinCtx(ctx context.Context, stdin io.Reader, pro
200204
IsError: false,
201205
}
202206
if err := applyCompletionHooks(ctx, opts, result); err != nil {
203-
return nil, err
207+
return result, err
204208
}
205209
return result, nil
206210
}

‎pkg/claude/execution_hooks.go‎

Lines changed: 4 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,6 @@ package claude
33
import (
44
"context"
55
"encoding/json"
6-
"errors"
76
"fmt"
87
)
98

@@ -83,7 +82,7 @@ func applyCompletionHooks(ctx context.Context, opts *RunOptions, result *ClaudeR
8382
}
8483

8584
if tracker := opts.BudgetTracker; tracker != nil && result.CostUSD > 0 {
86-
if err := tracker.AddSpend(result.SessionID, result.CostUSD); err != nil && !errors.Is(err, ErrBudgetExceeded) {
85+
if err := tracker.AddSpend(result.SessionID, result.CostUSD); err != nil {
8786
return err
8887
}
8988
}
@@ -97,26 +96,12 @@ func applyCompletionHooks(ctx context.Context, opts *RunOptions, result *ClaudeR
9796
return nil
9897
}
9998

100-
func preparePluginManager(ctx context.Context, pm *PluginManager) (func(), error) {
99+
func ensurePluginManagerInitialized(ctx context.Context, pm *PluginManager) error {
101100
if pm == nil {
102-
return func() {}, nil
103-
}
104-
105-
pm.mu.RLock()
106-
alreadyInitialized := pm.initialized
107-
pm.mu.RUnlock()
108-
109-
if alreadyInitialized {
110-
return func() {}, nil
111-
}
112-
113-
if err := pm.Initialize(ctx); err != nil {
114-
return nil, err
101+
return nil
115102
}
116103

117-
return func() {
118-
_ = pm.Shutdown(ctx)
119-
}, nil
104+
return pm.Initialize(ctx)
120105
}
121106

122107
type toolUseCall struct {

‎pkg/claude/execution_hooks_test.go‎

Lines changed: 171 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -2,9 +2,11 @@ package claude
22

33
import (
44
"context"
5+
"errors"
56
"os"
67
"os/exec"
78
"path/filepath"
9+
"strings"
810
"sync"
911
"testing"
1012
)
@@ -66,8 +68,9 @@ func TestRunPromptCtx_AppliesLifecycleHooks(t *testing.T) {
6668
execCommand = originalExecCommand
6769
}()
6870

69-
jsonOutput := `[{"type":"assistant","message":{"role":"assistant","content":[{"type":"tool_use","name":"Bash","input":{"command":"pwd"}}]},"session_id":"hook-session"},{"type":"result","subtype":"success","total_cost_usd":0.25,"duration_ms":12,"duration_api_ms":8,"is_error":false,"num_turns":1,"result":"done","session_id":"hook-session"}]`
70-
execCommand = mockExecCommandContext(t, []string{"-p", "Hook test", "--output-format", "json"}, jsonOutput, 0)
71+
jsonOutput := `{"type":"assistant","message":{"role":"assistant","content":[{"type":"tool_use","name":"Bash","input":{"command":"pwd"}}]},"session_id":"hook-session"}` + "\n" +
72+
`{"type":"result","subtype":"success","total_cost_usd":0.25,"duration_ms":12,"duration_api_ms":8,"is_error":false,"num_turns":1,"result":"done","session_id":"hook-session"}` + "\n"
73+
execCommand = mockExecCommandContext(t, []string{"-p", "Hook test", "--output-format", "stream-json", "--verbose"}, jsonOutput, 0)
7174

7275
pm := NewPluginManager()
7376
plugin := &lifecyclePlugin{}
@@ -103,8 +106,8 @@ func TestRunPromptCtx_AppliesLifecycleHooks(t *testing.T) {
103106
if plugin.initCount != 1 {
104107
t.Fatalf("Initialize count = %d, want 1", plugin.initCount)
105108
}
106-
if plugin.shutdownCount != 1 {
107-
t.Fatalf("Shutdown count = %d, want 1", plugin.shutdownCount)
109+
if plugin.shutdownCount != 0 {
110+
t.Fatalf("Shutdown count = %d, want 0", plugin.shutdownCount)
108111
}
109112
if plugin.messageCount != 2 {
110113
t.Fatalf("Message count = %d, want 2", plugin.messageCount)
@@ -129,8 +132,8 @@ func TestRunPromptCtx_DoesNotShutdownPreinitializedPluginManager(t *testing.T) {
129132
execCommand = originalExecCommand
130133
}()
131134

132-
jsonOutput := `{"type":"result","subtype":"success","total_cost_usd":0.01,"duration_ms":5,"duration_api_ms":5,"is_error":false,"num_turns":1,"result":"ok","session_id":"sticky-session"}`
133-
execCommand = mockExecCommandContext(t, []string{"-p", "Sticky hooks", "--output-format", "json"}, jsonOutput, 0)
135+
jsonOutput := `{"type":"result","subtype":"success","total_cost_usd":0.01,"duration_ms":5,"duration_api_ms":5,"is_error":false,"num_turns":1,"result":"ok","session_id":"sticky-session"}` + "\n"
136+
execCommand = mockExecCommandContext(t, []string{"-p", "Sticky hooks", "--output-format", "stream-json", "--verbose"}, jsonOutput, 0)
134137

135138
pm := NewPluginManager()
136139
plugin := &lifecyclePlugin{}
@@ -163,6 +166,123 @@ func TestRunPromptCtx_DoesNotShutdownPreinitializedPluginManager(t *testing.T) {
163166
}
164167
}
165168

169+
func TestRunPromptCtx_ReusesLazyInitializedPluginManager(t *testing.T) {
170+
originalExecCommand := execCommand
171+
defer func() {
172+
execCommand = originalExecCommand
173+
}()
174+
175+
jsonOutput := `{"type":"result","subtype":"success","total_cost_usd":0.01,"duration_ms":5,"duration_api_ms":5,"is_error":false,"num_turns":1,"result":"ok","session_id":"sticky-session"}` + "\n"
176+
execCommand = mockExecCommandContext(t, []string{"-p", "Sticky hooks", "--output-format", "stream-json", "--verbose"}, jsonOutput, 0)
177+
178+
pm := NewPluginManager()
179+
plugin := &lifecyclePlugin{}
180+
if err := pm.Register(plugin, nil); err != nil {
181+
t.Fatalf("Register() error = %v", err)
182+
}
183+
184+
client := NewClient("claude")
185+
for i := 0; i < 2; i++ {
186+
if _, err := client.RunPromptCtx(context.Background(), "Sticky hooks", &RunOptions{
187+
Format: JSONOutput,
188+
PluginManager: pm,
189+
}); err != nil {
190+
t.Fatalf("RunPromptCtx() error = %v", err)
191+
}
192+
}
193+
194+
plugin.mu.Lock()
195+
defer plugin.mu.Unlock()
196+
197+
if plugin.initCount != 1 {
198+
t.Fatalf("Initialize count = %d, want 1", plugin.initCount)
199+
}
200+
if plugin.shutdownCount != 0 {
201+
t.Fatalf("Shutdown count = %d, want 0", plugin.shutdownCount)
202+
}
203+
}
204+
205+
func TestRunFromStdinCtx_AppliesLifecycleHooks(t *testing.T) {
206+
originalExecCommand := execCommand
207+
defer func() {
208+
execCommand = originalExecCommand
209+
}()
210+
211+
jsonOutput := `{"type":"assistant","message":{"role":"assistant","content":[{"type":"tool_use","name":"Read","input":{"file_path":"README.md"}}]},"session_id":"stdin-hook-session"}` + "\n" +
212+
`{"type":"result","subtype":"success","total_cost_usd":0.05,"duration_ms":7,"duration_api_ms":5,"is_error":false,"num_turns":1,"result":"stdin-done","session_id":"stdin-hook-session"}` + "\n"
213+
execCommand = mockExecCommandContext(t, []string{"-p", "Hook stdin", "--output-format", "stream-json", "--verbose"}, jsonOutput, 0)
214+
215+
pm := NewPluginManager()
216+
plugin := &lifecyclePlugin{}
217+
if err := pm.Register(plugin, nil); err != nil {
218+
t.Fatalf("Register() error = %v", err)
219+
}
220+
221+
client := NewClient("claude")
222+
result, err := client.RunFromStdinCtx(context.Background(), strings.NewReader("stdin body"), "Hook stdin", &RunOptions{
223+
Format: JSONOutput,
224+
PluginManager: pm,
225+
})
226+
if err != nil {
227+
t.Fatalf("RunFromStdinCtx() error = %v", err)
228+
}
229+
if result.Result != "stdin-done" {
230+
t.Fatalf("Result = %q, want %q", result.Result, "stdin-done")
231+
}
232+
233+
plugin.mu.Lock()
234+
defer plugin.mu.Unlock()
235+
236+
if plugin.initCount != 1 {
237+
t.Fatalf("Initialize count = %d, want 1", plugin.initCount)
238+
}
239+
if plugin.shutdownCount != 0 {
240+
t.Fatalf("Shutdown count = %d, want 0", plugin.shutdownCount)
241+
}
242+
if plugin.messageCount != 2 {
243+
t.Fatalf("Message count = %d, want 2", plugin.messageCount)
244+
}
245+
if plugin.completeCount != 1 {
246+
t.Fatalf("Complete count = %d, want 1", plugin.completeCount)
247+
}
248+
if len(plugin.toolCalls) != 1 || plugin.toolCalls[0] != "Read" {
249+
t.Fatalf("Tool calls = %v, want [Read]", plugin.toolCalls)
250+
}
251+
if plugin.lastToolInput.FilePath != "README.md" {
252+
t.Fatalf("Tool file path = %q, want %q", plugin.lastToolInput.FilePath, "README.md")
253+
}
254+
}
255+
256+
func TestRunPromptCtx_BudgetExceededReturnsResultAndError(t *testing.T) {
257+
originalExecCommand := execCommand
258+
defer func() {
259+
execCommand = originalExecCommand
260+
}()
261+
262+
jsonOutput := `{"type":"result","subtype":"success","total_cost_usd":0.25,"duration_ms":12,"duration_api_ms":8,"is_error":false,"num_turns":1,"result":"done","session_id":"budget-session"}`
263+
execCommand = mockExecCommandContext(t, []string{"-p", "Budget test", "--output-format", "json"}, jsonOutput, 0)
264+
265+
tracker := NewBudgetTracker(&BudgetConfig{MaxBudgetUSD: 0.10})
266+
client := NewClient("claude")
267+
268+
result, err := client.RunPromptCtx(context.Background(), "Budget test", &RunOptions{
269+
Format: JSONOutput,
270+
BudgetTracker: tracker,
271+
})
272+
if !errors.Is(err, ErrBudgetExceeded) {
273+
t.Fatalf("RunPromptCtx() error = %v, want ErrBudgetExceeded", err)
274+
}
275+
if result == nil {
276+
t.Fatal("RunPromptCtx() result = nil, want result")
277+
}
278+
if result.Result != "done" {
279+
t.Fatalf("Result = %q, want %q", result.Result, "done")
280+
}
281+
if tracker.TotalSpent() != 0.25 {
282+
t.Fatalf("TotalSpent = %f, want %f", tracker.TotalSpent(), 0.25)
283+
}
284+
}
285+
166286
func TestStreamPrompt_AppliesLifecycleHooks(t *testing.T) {
167287
originalExecCommand := execCommand
168288
defer func() {
@@ -210,8 +330,8 @@ func TestStreamPrompt_AppliesLifecycleHooks(t *testing.T) {
210330
if plugin.initCount != 1 {
211331
t.Fatalf("Initialize count = %d, want 1", plugin.initCount)
212332
}
213-
if plugin.shutdownCount != 1 {
214-
t.Fatalf("Shutdown count = %d, want 1", plugin.shutdownCount)
333+
if plugin.shutdownCount != 0 {
334+
t.Fatalf("Shutdown count = %d, want 0", plugin.shutdownCount)
215335
}
216336
if plugin.messageCount != 2 {
217337
t.Fatalf("Message count = %d, want 2", plugin.messageCount)
@@ -227,6 +347,49 @@ func TestStreamPrompt_AppliesLifecycleHooks(t *testing.T) {
227347
}
228348
}
229349

350+
func TestStreamPrompt_BudgetExceededReturnsResultBeforeError(t *testing.T) {
351+
originalExecCommand := execCommand
352+
defer func() {
353+
execCommand = originalExecCommand
354+
}()
355+
356+
mockBinary := buildStreamingMockBinary(t)
357+
execCommand = func(ctx context.Context, name string, arg ...string) *exec.Cmd {
358+
return exec.Command(mockBinary)
359+
}
360+
361+
tracker := NewBudgetTracker(&BudgetConfig{MaxBudgetUSD: 0.10})
362+
client := NewClient("claude")
363+
messageCh, errCh := client.StreamPrompt(context.Background(), "Stream hooks", &RunOptions{
364+
BudgetTracker: tracker,
365+
})
366+
367+
var gotMessages []Message
368+
for msg := range messageCh {
369+
gotMessages = append(gotMessages, msg)
370+
}
371+
372+
var gotErr error
373+
for err := range errCh {
374+
if err != nil {
375+
gotErr = err
376+
}
377+
}
378+
379+
if !errors.Is(gotErr, ErrBudgetExceeded) {
380+
t.Fatalf("StreamPrompt() error = %v, want ErrBudgetExceeded", gotErr)
381+
}
382+
if len(gotMessages) != 2 {
383+
t.Fatalf("Expected 2 streamed messages, got %d", len(gotMessages))
384+
}
385+
if gotMessages[len(gotMessages)-1].Type != "result" {
386+
t.Fatalf("Last streamed message type = %q, want %q", gotMessages[len(gotMessages)-1].Type, "result")
387+
}
388+
if tracker.TotalSpent() != 0.4 {
389+
t.Fatalf("TotalSpent = %f, want %f", tracker.TotalSpent(), 0.4)
390+
}
391+
}
392+
230393
func buildStreamingMockBinary(t *testing.T) string {
231394
t.Helper()
232395

0 commit comments

Comments
 (0)