| 1 | package agent |
| 2 | |
| 3 | import ( |
| 4 | "context" |
| 5 | "encoding/json" |
| 6 | "fmt" |
| 7 | "io" |
| 8 | "os" |
| 9 | "path/filepath" |
| 10 | |
| 11 | "reasonix/internal/evidence" |
| 12 | "reasonix/internal/provider" |
| 13 | ) |
| 14 | |
| 15 | // ForkBundle freezes the full turn state at a policy's first eligibility so a |
| 16 | // control and a treatment continuation can start from the identical point. |
| 17 | // Versioned from day one: this format is shared infrastructure for every |
| 18 | // policy experiment (EBM, reasoning governor, delegation admission, rollback). |
| 19 | type ForkBundle struct { |
| 20 | Version int `json:"version"` |
| 21 | Policy string `json:"policy"` |
| 22 | Input string `json:"input"` |
| 23 | EligibleRound int `json:"eligible_round"` |
| 24 | BlindAtFork int `json:"blind_at_fork"` |
| 25 | DebtAtFork int `json:"debt_at_fork"` |
| 26 | MutatedBases []string `json:"mutated_bases,omitempty"` |
| 27 | LocalExecSeen bool `json:"local_exec_seen,omitempty"` |
| 28 | RunwayBalance int `json:"runway_balance,omitempty"` |
| 29 | RunwayDry int `json:"runway_dry,omitempty"` |
| 30 | RunwayIdle int `json:"runway_idle,omitempty"` |
| 31 | RunwayObserved bool `json:"runway_observed,omitempty"` |
| 32 | Messages []provider.Message `json:"messages"` |
| 33 | } |
| 34 | |
| 35 | const forkBundleVersion = 1 |
| 36 | |
| 37 | // govReasoningThreshold marks a round's thinking as expensive enough that a |
| 38 | // governor experiment wants the state frozen before the next purchase. |
| 39 | |
| 40 | // forkCapturePolicy selects which policy's trigger owns bundle capture; |
| 41 | // unset defaults to the EBM trigger for compatibility with existing runs. |
| 42 | func forkCapturePolicy() string { |
| 43 | if os.Getenv("REASONIX_EXPERIMENT_FORK_CAPTURE_DIR") == "" { |
| 44 | return "" |
| 45 | } |
| 46 | if p := os.Getenv("REASONIX_EXPERIMENT_FORK_POLICY"); p != "" { |
| 47 | return p |
| 48 | } |
| 49 | return "ebm" |
| 50 | } |
| 51 | |
| 52 | // forkCaptureProvider snapshots the session at the next Stream call after |
| 53 | // eligibility was armed — the exact state the uninterrupted run sends. |
| 54 | type forkCaptureProvider struct { |
| 55 | inner provider.Provider |
| 56 | a *Agent |
| 57 | } |
| 58 | |
| 59 | func (p *forkCaptureProvider) ReasoningCapability() provider.ReasoningCapability { |
| 60 | if owner, ok := p.inner.(provider.ReasoningProvider); ok { |
| 61 | return owner.ReasoningCapability() |
| 62 | } |
| 63 | return provider.ReasoningOptions("") |
| 64 | } |
| 65 | |
| 66 | func (p *forkCaptureProvider) Name() string { return p.inner.Name() } |
| 67 | |
| 68 | func (p *forkCaptureProvider) ModelInfo() provider.ModelInfo { |
| 69 | if info, ok := p.inner.(provider.ModelInfoProvider); ok { |
| 70 | return info.ModelInfo() |
| 71 | } |
| 72 | return provider.ModelInfo{} |
| 73 | } |
| 74 | |
| 75 | func (p *forkCaptureProvider) OutputBudget() int { return outputBudgetOf(p.inner) } |
| 76 | |
| 77 | func (p *forkCaptureProvider) SharesContextWindow() bool { return sharesContextWindow(p.inner) } |
| 78 | |
| 79 | func (p *forkCaptureProvider) ContextBudgetPolicy() provider.ContextBudgetPolicy { |
| 80 | return provider.ResolveContextBudgetPolicy(p.inner) |
| 81 | } |
| 82 | |
| 83 | func (p *forkCaptureProvider) SharedWindowInputPolicy() provider.SharedWindowInputPolicy { |
| 84 | return sharedWindowInputPolicyOf(p.inner) |
| 85 | } |
| 86 | |
| 87 | func (p *forkCaptureProvider) Stream(ctx context.Context, req provider.Request) (<-chan provider.Chunk, error) { |
| 88 | a := p.a |
| 89 | if a.task.ebm.captureArmed && !a.task.ebm.captured { |
| 90 | a.task.ebm.captured = true |
| 91 | messages := a.sess.conversation.Snapshot() |
| 92 | seed := a.task.outcome.ForkSeed() |
| 93 | b := ForkBundle{ |
| 94 | Version: forkBundleVersion, Policy: forkCapturePolicy(), |
| 95 | Input: forkTurnInput(messages), |
| 96 | EligibleRound: a.task.ebm.captureRound, BlindAtFork: seed.BlindMutations, |
| 97 | DebtAtFork: seed.DebtAge, MutatedBases: seed.MutatedBases, |
| 98 | LocalExecSeen: seed.LocalExecSeen, |
| 99 | RunwayBalance: seed.RunwayBalance, RunwayDry: seed.RunwayDry, |
| 100 | RunwayIdle: seed.RunwayIdle, RunwayObserved: seed.RunwayObserved, |
| 101 | Messages: messages, |
| 102 | } |
| 103 | if err := writeForkBundle(os.Getenv("REASONIX_EXPERIMENT_FORK_CAPTURE_DIR"), b); err != nil { |
| 104 | fmt.Fprintln(os.Stderr, "fork capture:", err) |
| 105 | } |
| 106 | } |
| 107 | return p.inner.Stream(ctx, req) |
| 108 | } |
| 109 | |
| 110 | // forkTurnInput recovers the turn's raw input from the frozen conversation: |
| 111 | // the first user message's authored form. Single-turn scope (the bench runs |
| 112 | // one turn per task); multi-turn capture would need the active turn's index. |
| 113 | func forkTurnInput(messages []provider.Message) string { |
| 114 | for _, m := range messages { |
| 115 | if IsUserAuthoredTurnMessage(m) { |
| 116 | if m.RawContent != "" { |
| 117 | return m.RawContent |
| 118 | } |
| 119 | return m.Content |
| 120 | } |
| 121 | } |
| 122 | return "" |
| 123 | } |
| 124 | |
| 125 | func writeForkBundle(dir string, b ForkBundle) error { |
| 126 | if err := os.MkdirAll(dir, 0o755); err != nil { |
| 127 | return err |
| 128 | } |
| 129 | data, err := json.Marshal(b) |
| 130 | if err != nil { |
| 131 | return err |
| 132 | } |
| 133 | if err := os.WriteFile(filepath.Join(dir, "bundle.json"), data, 0o644); err != nil { |
| 134 | return err |
| 135 | } |
| 136 | cwd, err := os.Getwd() |
| 137 | if err != nil { |
| 138 | return err |
| 139 | } |
| 140 | return copyWorkspace(cwd, filepath.Join(dir, "workspace")) |
| 141 | } |
| 142 | |
| 143 | // copyWorkspace mirrors the task workdir minus harness artifacts, so a |
| 144 | // continuation starts from byte-identical files. |
| 145 | func copyWorkspace(src, dst string) error { |
| 146 | return filepath.WalkDir(src, func(path string, d os.DirEntry, err error) error { |
| 147 | if err != nil { |
| 148 | return err |
| 149 | } |
| 150 | rel, rerr := filepath.Rel(src, path) |
| 151 | if rerr != nil || rel == "." { |
| 152 | return rerr |
| 153 | } |
| 154 | name := d.Name() |
| 155 | if d.IsDir() { |
| 156 | if name == "__pycache__" || name == ".git" { |
| 157 | return filepath.SkipDir |
| 158 | } |
| 159 | return os.MkdirAll(filepath.Join(dst, rel), 0o755) |
| 160 | } |
| 161 | if name == ".run-metrics.json" { |
| 162 | return nil |
| 163 | } |
| 164 | in, oerr := os.Open(path) |
| 165 | if oerr != nil { |
| 166 | return oerr |
| 167 | } |
| 168 | defer in.Close() |
| 169 | if err := os.MkdirAll(filepath.Dir(filepath.Join(dst, rel)), 0o755); err != nil { |
| 170 | return err |
| 171 | } |
| 172 | out, cerr := os.Create(filepath.Join(dst, rel)) |
| 173 | if cerr != nil { |
| 174 | return cerr |
| 175 | } |
| 176 | defer out.Close() |
| 177 | _, err = io.Copy(out, in) |
| 178 | return err |
| 179 | }) |
| 180 | } |
| 181 | |
| 182 | // LoadForkBundle reads and version-checks a bundle. |
| 183 | func LoadForkBundle(path string) (*ForkBundle, error) { |
| 184 | data, err := os.ReadFile(path) |
| 185 | if err != nil { |
| 186 | return nil, err |
| 187 | } |
| 188 | var b ForkBundle |
| 189 | if err := json.Unmarshal(data, &b); err != nil { |
| 190 | return nil, err |
| 191 | } |
| 192 | if b.Version != forkBundleVersion { |
| 193 | return nil, fmt.Errorf("fork bundle version %d, this build replays %d", b.Version, forkBundleVersion) |
| 194 | } |
| 195 | return &b, nil |
| 196 | } |
| 197 | |
| 198 | // actFirstNudge is the reasoning-governor's soft treatment: shaping, not |
| 199 | // capping — spend cheap external evidence before expensive speculation. |
| 200 | const actFirstNudge = "[guidance] Prefer cheap repository evidence or a targeted check over extended " + |
| 201 | "speculation when either can reduce uncertainty." |
| 202 | |
| 203 | // armForkContinuation makes the next Run continue from the bundle: turn-local |
| 204 | // classification still runs on the same input, then the appended user message |
| 205 | // is replaced wholesale by the frozen conversation. A non-empty nudge is the |
| 206 | // arm's single treatment, placed in the live policy's slot; the dose disarms |
| 207 | // every runtime policy for the continuation. |
| 208 | func (a *Agent) armForkContinuation(b *ForkBundle, nudge string) { |
| 209 | a.pending.forkRestore = func(_ *turnRuntime) { |
| 210 | messages := append([]provider.Message(nil), b.Messages...) |
| 211 | if nudge != "" { |
| 212 | applyForkTreatment(messages, nudge) |
| 213 | } |
| 214 | a.sess.conversation.Replace(messages) |
| 215 | a.task.outcome = evidence.RestoreOutcomeTracker(evidence.OutcomeSeed{ |
| 216 | MutatedBases: b.MutatedBases, DebtAge: b.DebtAtFork, |
| 217 | BlindMutations: b.BlindAtFork, LocalExecSeen: b.LocalExecSeen, |
| 218 | RunwayBalance: b.RunwayBalance, RunwayDry: b.RunwayDry, |
| 219 | RunwayIdle: b.RunwayIdle, RunwayObserved: b.RunwayObserved, |
| 220 | }) |
| 221 | a.task.ebm = ebmState{fired: true, captured: true, captureRound: b.EligibleRound} |
| 222 | } |
| 223 | } |
| 224 | |
| 225 | // applyForkTreatment appends the nudge to the eligible batch's first tool |
| 226 | // result — the same slot the live policy writes to. |
| 227 | func applyForkTreatment(messages []provider.Message, nudge string) { |
| 228 | lastAssistant := -1 |
| 229 | for i, m := range messages { |
| 230 | if m.Role == provider.RoleAssistant && len(m.ToolCalls) > 0 { |
| 231 | lastAssistant = i |
| 232 | } |
| 233 | } |
| 234 | for i := lastAssistant + 1; lastAssistant >= 0 && i < len(messages); i++ { |
| 235 | if messages[i].Role == provider.RoleTool { |
| 236 | messages[i].Content += "\n\n" + nudge |
| 237 | return |
| 238 | } |
| 239 | } |
| 240 | } |
| 241 | |
| 242 | // maybeWrapForkCaptureProvider interposes the capture wrapper when the |
| 243 | // experiment env asks for bundles; inert otherwise. |
| 244 | func (a *Agent) maybeWrapForkCaptureProvider() { |
| 245 | if os.Getenv("REASONIX_EXPERIMENT_FORK_CAPTURE_DIR") != "" && a.svc.prov != nil { |
| 246 | a.svc.prov = &forkCaptureProvider{inner: a.svc.prov, a: a} |
| 247 | } |
| 248 | } |
| 249 | |
| 250 | // maybeArmForkFromEnv wires the experiment from the environment so the bench |
| 251 | // can fork without new public plumbing. Control is the default arm. |
| 252 | func (a *Agent) maybeArmForkFromEnv() { |
| 253 | path := os.Getenv("REASONIX_EXPERIMENT_FORK_BUNDLE") |
| 254 | if path == "" { |
| 255 | return |
| 256 | } |
| 257 | b, err := LoadForkBundle(path) |
| 258 | if err != nil { |
| 259 | fmt.Fprintln(os.Stderr, "fork continuation:", err) |
| 260 | return |
| 261 | } |
| 262 | nudge := "" |
| 263 | switch os.Getenv("REASONIX_EXPERIMENT_FORK_ARM") { |
| 264 | case "actfirst": |
| 265 | nudge = actFirstNudge |
| 266 | } |
| 267 | a.armForkContinuation(b, nudge) |
| 268 | } |
| 269 |