From 03a196e613b8b39d89d0ffccd387a12fc3fd43f2 Mon Sep 17 00:00:00 2001 From: ntduc Date: Wed, 22 Jul 2026 14:13:23 +0700 Subject: [PATCH] feat(hooks): add script context injection --- internal/hooks/dispatcher.go | 6 +++++ internal/hooks/dispatcher_test.go | 31 ++++++++++++++++++++++- internal/hooks/handlers/script.go | 3 ++- internal/hooks/handlers/script_runtime.go | 23 +++++++++++------ internal/hooks/handlers/script_test.go | 12 +++++++++ internal/hooks/script_result.go | 6 ++++- internal/hooks/types.go | 14 +++++----- internal/pipeline/context_stage.go | 12 +++++++-- internal/pipeline/stages_test.go | 26 +++++++++++++++++++ 9 files changed, 113 insertions(+), 20 deletions(-) diff --git a/internal/hooks/dispatcher.go b/internal/hooks/dispatcher.go index 6fe6ca21..35287eea 100644 --- a/internal/hooks/dispatcher.go +++ b/internal/hooks/dispatcher.go @@ -7,6 +7,7 @@ import ( "fmt" "log/slog" "maps" + "strings" "sync" "sync/atomic" "time" @@ -198,6 +199,7 @@ func (d *stdDispatcher) runSync(ctx context.Context, ev Event, chain []HookConfi evMut.ToolInput = cloneMap(ev.ToolInput) } mutated := false + additionalContexts := make([]string, 0, len(chain)) for _, cfg := range chain { if !cfg.Enabled { @@ -259,6 +261,9 @@ func (d *stdDispatcher) runSync(ctx context.Context, ev Event, chain []HookConfi ) } } + if dec == DecisionAllow && scriptRes.AdditionalContext != "" { + additionalContexts = append(additionalContexts, scriptRes.AdditionalContext) + } switch dec { case DecisionBlock: @@ -282,6 +287,7 @@ func (d *stdDispatcher) runSync(ctx context.Context, ev Event, chain []HookConfi } result := FireResult{Decision: DecisionAllow} + result.AdditionalContext = strings.Join(additionalContexts, "\n\n") if mutated { if evMut.ToolInput != nil { result.UpdatedToolInput = evMut.ToolInput diff --git a/internal/hooks/dispatcher_test.go b/internal/hooks/dispatcher_test.go index 35dd0fa5..a2d01b7a 100644 --- a/internal/hooks/dispatcher_test.go +++ b/internal/hooks/dispatcher_test.go @@ -411,16 +411,45 @@ func TestNoopDispatcher_AlwaysAllows(t *testing.T) { // the dispatcher's source-tier gate can be exercised without invoking the // real goja runtime. type mutatingHandler struct { - updated map[string]any + updated map[string]any + additionalContext string } func (h *mutatingHandler) Execute(ctx context.Context, _ hooks.HookConfig, _ hooks.Event) (hooks.Decision, error) { if r := hooks.ScriptResultFrom(ctx); r != nil { r.UpdatedInput = h.updated + r.AdditionalContext = h.additionalContext } return hooks.DecisionAllow, nil } +func TestDispatcher_UIScriptAdditionalContextAllowedWithoutInputMutation(t *testing.T) { + cfg := newBaseHook(hooks.HandlerScript, hooks.EventUserPromptSubmit) + cfg.Source = "ui" + fs := &fakeStore{hooks: []hooks.HookConfig{cfg}} + h := &mutatingHandler{ + updated: map[string]any{"rawInput": "must-not-apply"}, + additionalContext: "Call the required workflow tool.", + } + d := hooks.NewStdDispatcher(hooks.StdDispatcherOpts{ + Store: fs, Audit: hooks.NewAuditWriter(fs, ""), + Handlers: map[hooks.HandlerType]hooks.Handler{hooks.HandlerScript: h}, + }) + + r, err := d.Fire(context.Background(), hooks.Event{ + EventID: "e-context", HookEvent: hooks.EventUserPromptSubmit, RawInput: "original", + }) + if err != nil || r.Decision != hooks.DecisionAllow { + t.Fatalf("decision=%q err=%v", r.Decision, err) + } + if r.AdditionalContext != "Call the required workflow tool." { + t.Fatalf("additional context=%q", r.AdditionalContext) + } + if r.UpdatedRawInput != nil { + t.Fatalf("UI source mutated raw input: %q", *r.UpdatedRawInput) + } +} + // TestDispatcher_ScriptMutation_BuiltinSourceApplies verifies the source-tier // capability: a script hook with source="builtin" has its UpdatedInput applied // to FireResult.Updated*; a script hook with source="ui" does NOT. diff --git a/internal/hooks/handlers/script.go b/internal/hooks/handlers/script.go index 04aa9d9e..6e4cb066 100644 --- a/internal/hooks/handlers/script.go +++ b/internal/hooks/handlers/script.go @@ -187,7 +187,7 @@ func (h *ScriptHandler) Execute(ctx context.Context, cfg hooks.HookConfig, ev ho return hooks.DecisionError, sanitizeError(execErr) } - dec, reason, updated, parseErr := parseReturn(rt, result) + dec, reason, additionalContext, updated, parseErr := parseReturn(rt, result) if parseErr != nil { return hooks.DecisionError, parseErr } @@ -208,6 +208,7 @@ func (h *ScriptHandler) Execute(ctx context.Context, cfg hooks.HookConfig, ev ho } if r := hooks.ScriptResultFrom(ctx); r != nil { r.Reason = reason + r.AdditionalContext = additionalContext r.UpdatedInput = updated r.Stdout = stdout.String() } diff --git a/internal/hooks/handlers/script_runtime.go b/internal/hooks/handlers/script_runtime.go index 0d9becf5..57d2b928 100644 --- a/internal/hooks/handlers/script_runtime.go +++ b/internal/hooks/handlers/script_runtime.go @@ -134,31 +134,32 @@ func captureStdout(rt *goja.Runtime, buf *strings.Builder) { _ = rt.Set("console", console) } -// parseReturn enforces the `{decision, reason?, updatedInput?}` contract. +// parseReturn enforces the +// `{decision, reason?, additionalContext?, updatedInput?}` contract. // Any deviation (non-object, missing field, unknown decision string) maps to // DecisionError with a descriptive error. -func parseReturn(rt *goja.Runtime, v goja.Value) (hooks.Decision, string, map[string]any, error) { +func parseReturn(rt *goja.Runtime, v goja.Value) (hooks.Decision, string, string, map[string]any, error) { if v == nil || goja.IsUndefined(v) || goja.IsNull(v) { - return hooks.DecisionError, "", nil, errors.New("script return: must be object with decision") + return hooks.DecisionError, "", "", nil, errors.New("script return: must be object with decision") } obj := v.ToObject(rt) if obj == nil { - return hooks.DecisionError, "", nil, errors.New("script return: must be object") + return hooks.DecisionError, "", "", nil, errors.New("script return: must be object") } decRaw := obj.Get("decision") if decRaw == nil || goja.IsUndefined(decRaw) { - return hooks.DecisionError, "", nil, errors.New("script return: missing decision field") + return hooks.DecisionError, "", "", nil, errors.New("script return: missing decision field") } decStr, ok := decRaw.Export().(string) if !ok { - return hooks.DecisionError, "", nil, errors.New("script return: decision must be string") + return hooks.DecisionError, "", "", nil, errors.New("script return: decision must be string") } dec := hooks.Decision(decStr) switch dec { case hooks.DecisionAllow, hooks.DecisionBlock, hooks.DecisionAsk, hooks.DecisionDefer: // ok default: - return hooks.DecisionError, "", nil, fmt.Errorf("script return: invalid decision %q", decStr) + return hooks.DecisionError, "", "", nil, fmt.Errorf("script return: invalid decision %q", decStr) } var reason string if r := obj.Get("reason"); r != nil && !goja.IsUndefined(r) { @@ -166,11 +167,17 @@ func parseReturn(rt *goja.Runtime, v goja.Value) (hooks.Decision, string, map[st reason = s } } + var additionalContext string + if a := obj.Get("additionalContext"); a != nil && !goja.IsUndefined(a) { + if s, ok := a.Export().(string); ok { + additionalContext = s + } + } var updated map[string]any if u := obj.Get("updatedInput"); u != nil && !goja.IsUndefined(u) { if m, ok := u.Export().(map[string]any); ok { updated = m } } - return dec, reason, updated, nil + return dec, reason, additionalContext, updated, nil } diff --git a/internal/hooks/handlers/script_test.go b/internal/hooks/handlers/script_test.go index 1a983c9e..acca49b1 100644 --- a/internal/hooks/handlers/script_test.go +++ b/internal/hooks/handlers/script_test.go @@ -73,6 +73,18 @@ func TestHappyPathAllow(t *testing.T) { } } +func TestAllowReturnsAdditionalContext(t *testing.T) { + src := `function handle(event) { return {decision: "allow", additionalContext: "Use the required workflow."}; }` + h := newTestHandler() + dec, err, res := runWithResult(t, h, mkCfg(src), mkEvent(), 500*time.Millisecond) + if err != nil || dec != hooks.DecisionAllow { + t.Fatalf("decision=%v err=%v", dec, err) + } + if res.AdditionalContext != "Use the required workflow." { + t.Fatalf("additional context=%q", res.AdditionalContext) + } +} + // TestHappyPathBlock verifies block + reason round-trips. func TestHappyPathBlock(t *testing.T) { src := `function handle(event) { return {decision: "block", reason: "nope"}; }` diff --git a/internal/hooks/script_result.go b/internal/hooks/script_result.go index 83725e6a..b0f7180c 100644 --- a/internal/hooks/script_result.go +++ b/internal/hooks/script_result.go @@ -3,7 +3,7 @@ package hooks import "context" // ScriptResult carries non-decision outputs from a script handler execution -// (reason text, updatedInput proposal, captured stdout). The dispatcher +// (reason text, additional context, updatedInput proposal, captured stdout). The dispatcher // provisions one via context before calling Handler.Execute; the handler // populates it on success. Phase 03 wires this into the dispatcher so that // only builtin-source hooks can apply UpdatedInput to the pipeline state. @@ -13,6 +13,10 @@ import "context" type ScriptResult struct { // Reason is the human-readable explanation returned by the script. Reason string + // AdditionalContext is instruction text appended to the current run's + // extra system prompt after the hook returns allow. It does not mutate the + // original user input. + AdditionalContext string // UpdatedInput is a proposed replacement for Event.ToolInput. Dispatcher // applies only when cfg.Source == "builtin" (source-tier capability). UpdatedInput map[string]any diff --git a/internal/hooks/types.go b/internal/hooks/types.go index 2dd55d20..8aedfeb5 100644 --- a/internal/hooks/types.go +++ b/internal/hooks/types.go @@ -151,12 +151,12 @@ func (d Decision) IsBlock() bool { // DecisionReason carries a human-readable explanation when Decision is // DecisionBlock (or DecisionAsk/DecisionDefer treated as block). This is // injected as a user message to trigger a retry iteration. -// type FireResult struct { - Decision Decision - DecisionReason string - UpdatedToolInput map[string]any - UpdatedRawInput *string + Decision Decision + DecisionReason string + AdditionalContext string + UpdatedToolInput map[string]any + UpdatedRawInput *string } // ─── Config & execution structs ────────────────────────────────────────────── @@ -233,7 +233,7 @@ type Event struct { HookEvent HookEvent // PostModelResponse fields (populated when HookEvent == EventPostModelResponse). - ModelResponse string // the generated response content - Thinking string // reasoning content (if any) + ModelResponse string // the generated response content + Thinking string // reasoning content (if any) ToolCalls []providers.ToolCall // empty for final response, populated if tool calls present } diff --git a/internal/pipeline/context_stage.go b/internal/pipeline/context_stage.go index 5f4d7050..04e709a3 100644 --- a/internal/pipeline/context_stage.go +++ b/internal/pipeline/context_stage.go @@ -60,8 +60,16 @@ func (s *ContextStage) Execute(ctx context.Context, state *RunState) error { HookEvent: hooks.EventUserPromptSubmit, }); r.Decision == hooks.DecisionBlock { return fmt.Errorf("hook blocked user_prompt_submit") - } else if r.UpdatedRawInput != nil { - state.Input.Message = *r.UpdatedRawInput + } else { + if r.UpdatedRawInput != nil { + state.Input.Message = *r.UpdatedRawInput + } + if r.AdditionalContext != "" { + if state.Input.ExtraSystemPrompt != "" { + state.Input.ExtraSystemPrompt += "\n\n" + } + state.Input.ExtraSystemPrompt += r.AdditionalContext + } } // 0. Inject context values (agent/tenant/user/workspace scoping, input guard, truncation). diff --git a/internal/pipeline/stages_test.go b/internal/pipeline/stages_test.go index 07ed597a..fb7b2960 100644 --- a/internal/pipeline/stages_test.go +++ b/internal/pipeline/stages_test.go @@ -2631,6 +2631,32 @@ func TestFinalizeStage_PopulatesFileSizes_SkipsNonexistent(t *testing.T) { // --- ContextStage tests --- +func TestContextStage_AppendsHookAdditionalContextToSystemPromptInput(t *testing.T) { + t.Parallel() + var capturedExtra string + deps := &PipelineDeps{ + Hooks: toolStageHookDispatcher{fire: func(_ context.Context, ev hooks.Event) (hooks.FireResult, error) { + if ev.HookEvent == hooks.EventUserPromptSubmit { + return hooks.FireResult{Decision: hooks.DecisionAllow, AdditionalContext: "mandatory workflow"}, nil + } + return hooks.FireResult{Decision: hooks.DecisionAllow}, nil + }}, + BuildMessages: func(_ context.Context, input *RunInput, _ []providers.Message, _ string) ([]providers.Message, error) { + capturedExtra = input.ExtraSystemPrompt + return []providers.Message{{Role: "system", Content: input.ExtraSystemPrompt}}, nil + }, + } + state := defaultState() + state.Input.ExtraSystemPrompt = "existing context" + stage := NewContextStage(deps) + if err := stage.Execute(context.Background(), state); err != nil { + t.Fatalf("Execute() error: %v", err) + } + if capturedExtra != "existing context\n\nmandatory workflow" { + t.Fatalf("extra system prompt=%q", capturedExtra) + } +} + func TestContextStage_ResolveWorkspace(t *testing.T) { t.Parallel() wantWS := &workspace.WorkspaceContext{ActivePath: "/resolved"}