返回 DeepSeek-Reasonix
probe.go
根目录 / internal / environment / probe.go
1 // Package environment probes the local developer environment at startup and
2 // renders a small, stable model-facing summary.
3 package environment
4
5 import (
6 "bytes"
7 "context"
8 "fmt"
9 "os"
10 "os/exec"
11 "path/filepath"
12 "runtime"
13 "sort"
14 "strings"
15 "sync"
16 "time"
17
18 "reasonix/internal/proc"
19 "reasonix/internal/secrets"
20 "reasonix/internal/shellparse"
21 )
22
23 const ProbeTimeout = 2 * time.Second
24
25 const probeWaitDelay = time.Second
26
27 const probeCacheTTL = 5 * time.Minute
28
29 const maxRenderedTools = 24
30
31 var probeTimeout = ProbeTimeout
32
33 type probeCacheEntry struct {
34 storedAt time.Time
35 results []ProbeResult
36 }
37
38 type probeInflight struct {
39 done chan struct{}
40 results []ProbeResult
41 }
42
43 var (
44 probeCacheMu sync.Mutex
45 probeCache = map[string]probeCacheEntry{}
46 probeInflightCalls = map[string]*probeInflight{}
47 probeNow = time.Now
48 )
49
50 type ProbeResult struct {
51 Command string
52 Binary string
53 Output string
54 Found bool
55 Error string
56 }
57
58 type ProbeOptions struct {
59 Overrides map[string]string
60 DenyRoots []string
61 // SnapshotDir, when set, persists probe results across process restarts
62 // (one snapshot per fingerprint under SnapshotDir/environment). A snapshot
63 // younger than probeSnapshotTTL is served without re-probing, and a
64 // refresh merges transient failures against the previous snapshot — both
65 // keep the rendered environment section byte-stable so the cached
66 // system-prompt prefix survives rebuilds. Empty disables persistence.
67 // The directory is host state, never model-visible, so it stays out of
68 // the probe fingerprint.
69 SnapshotDir string
70 }
71
72 func DefaultProbes() []string {
73 return []string{
74 "go version",
75 "python3 --version",
76 "python --version",
77 "node --version",
78 "npm --version",
79 "rustc --version",
80 "cargo --version",
81 "git version",
82 "make --version",
83 "rg --version",
84 "docker --version",
85 }
86 }
87
88 func RunProbes(ctx context.Context, commands []string) []ProbeResult {
89 return RunProbesWithOverrides(ctx, commands, nil)
90 }
91
92 func RunProbesWithOverrides(ctx context.Context, commands []string, overrides map[string]string) []ProbeResult {
93 return RunProbesWithOptions(ctx, commands, ProbeOptions{Overrides: overrides})
94 }
95
96 func RunProbesWithOptions(ctx context.Context, commands []string, opts ProbeOptions) []ProbeResult {
97 key := probeFingerprint(commands, opts)
98 now := probeNow()
99 if results, ok := cachedProbeResults(key, now); ok {
100 return results
101 }
102 if call, ok := beginProbe(key); ok {
103 <-call.done
104 return cloneProbeResults(call.results)
105 }
106 // A fresh persisted snapshot substitutes for a live run entirely: rebuilds
107 // and app relaunches within the TTL render the exact bytes the sessions on
108 // this machine were recorded with, so the provider prefix cache survives.
109 snapshot, hasSnapshot := loadProbeSnapshot(opts.SnapshotDir, key)
110 if hasSnapshot && now.Sub(snapshot.StoredAt) < probeSnapshotTTL {
111 finishProbe(key, snapshot.Results, now)
112 return cloneProbeResults(snapshot.Results)
113 }
114 results := runProbesUncached(ctx, commands, opts)
115 if hasSnapshot {
116 // Even an expired snapshot anchors the flap merge: transient failures
117 // (timeout, nonzero exit) keep the previous successful observation so
118 // a slow tool cannot rewrite the prompt prefix.
119 results = mergeProbeSnapshot(snapshot.Results, results)
120 }
121 saveProbeSnapshot(opts.SnapshotDir, key, results, now)
122 finishProbe(key, results, probeNow())
123 return cloneProbeResults(results)
124 }
125
126 func runProbesUncached(ctx context.Context, commands []string, opts ProbeOptions) []ProbeResult {
127 results := make([]ProbeResult, len(commands))
128 var wg sync.WaitGroup
129 for i, command := range commands {
130 wg.Add(1)
131 go func(i int, command string) {
132 defer wg.Done()
133 results[i] = runOne(ctx, command, opts)
134 }(i, command)
135 }
136 wg.Wait()
137 sortResults(results)
138 return results
139 }
140
141 func cachedProbeResults(key string, now time.Time) ([]ProbeResult, bool) {
142 probeCacheMu.Lock()
143 defer probeCacheMu.Unlock()
144 entry, ok := probeCache[key]
145 if !ok || now.Sub(entry.storedAt) >= probeCacheTTL {
146 if ok {
147 delete(probeCache, key)
148 }
149 return nil, false
150 }
151 return cloneProbeResults(entry.results), true
152 }
153
154 func beginProbe(key string) (*probeInflight, bool) {
155 probeCacheMu.Lock()
156 defer probeCacheMu.Unlock()
157 if call, ok := probeInflightCalls[key]; ok {
158 return call, true
159 }
160 probeInflightCalls[key] = &probeInflight{done: make(chan struct{})}
161 return nil, false
162 }
163
164 func finishProbe(key string, results []ProbeResult, now time.Time) {
165 probeCacheMu.Lock()
166 defer probeCacheMu.Unlock()
167 cached := cloneProbeResults(results)
168 probeCache[key] = probeCacheEntry{storedAt: now, results: cached}
169 if call, ok := probeInflightCalls[key]; ok {
170 call.results = cached
171 delete(probeInflightCalls, key)
172 close(call.done)
173 }
174 }
175
176 func probeFingerprint(commands []string, opts ProbeOptions) string {
177 var b strings.Builder
178 b.WriteString("v1")
179 for _, command := range commands {
180 b.WriteByte('\x00')
181 b.WriteString(strings.TrimSpace(command))
182 }
183 for _, name := range sortedMapKeys(opts.Overrides) {
184 b.WriteByte('\x00')
185 b.WriteString(name)
186 b.WriteByte('=')
187 b.WriteString(expandHome(opts.Overrides[name]))
188 }
189 for _, root := range normalizedDenyRoots(opts.DenyRoots) {
190 b.WriteByte('\x00')
191 b.WriteString("deny=")
192 b.WriteString(root)
193 }
194 return b.String()
195 }
196
197 func cloneProbeResults(results []ProbeResult) []ProbeResult {
198 if results == nil {
199 return nil
200 }
201 return append([]ProbeResult(nil), results...)
202 }
203
204 func runOne(ctx context.Context, command string, opts ProbeOptions) ProbeResult {
205 probe, err := shellparse.ParseStaticCommand(command, shellparse.StaticCommandPolicy{AllowEnvAssignments: true, AllowStderrToStdout: true})
206 if err != nil {
207 return ProbeResult{Command: command, Binary: command, Error: "invalid command: " + err.Error()}
208 }
209 parts := probe.Argv
210 if len(parts) == 0 {
211 return ProbeResult{Command: command, Binary: command, Error: "empty command"}
212 }
213 res := ProbeResult{Command: command, Binary: parts[0]}
214 var exe string
215 if override := strings.TrimSpace(opts.Overrides[parts[0]]); override != "" {
216 exe = expandHome(override)
217 if !filepath.IsAbs(exe) {
218 res.Error = "not trusted"
219 return res
220 }
221 if !fileExecutable(exe) {
222 res.Error = "not found"
223 return res
224 }
225 } else {
226 found, err := exec.LookPath(parts[0])
227 if err != nil {
228 res.Error = "not found"
229 return res
230 }
231 exe = found
232 }
233 if blockedExecutable(exe, opts.DenyRoots) {
234 res.Error = "not trusted"
235 return res
236 }
237 cmdCtx, cancel := context.WithTimeout(ctx, probeTimeout)
238 defer cancel()
239 cmd := exec.CommandContext(cmdCtx, exe, parts[1:]...)
240 // Always set the env explicitly: leaving cmd.Env nil would inherit the
241 // full process environment and bypass [secrets] filter_subprocess_env for
242 // probes that declare no extra variables of their own.
243 cmd.Env = append(secrets.ProcessEnv(), probe.Env...)
244 prepareProbeCommand(cmd)
245 var stdout, stderr bytes.Buffer
246 cmd.Stdout = &stdout
247 if probe.MergeStderr {
248 cmd.Stderr = &stdout
249 } else {
250 cmd.Stderr = &stderr
251 }
252 err = cmd.Run()
253 out := strings.TrimSpace(stdout.String())
254 if out == "" {
255 out = strings.TrimSpace(stderr.String())
256 }
257 if err != nil {
258 if cmdCtx.Err() == context.DeadlineExceeded {
259 res.Error = "timeout"
260 return res
261 }
262 if out == "" {
263 res.Error = "exit " + err.Error()
264 return res
265 }
266 res.Error = firstLine(out)
267 return res
268 }
269 res.Found = true
270 res.Output = firstLine(out)
271 return res
272 }
273
274 func prepareProbeCommand(cmd *exec.Cmd) {
275 proc.HideWindow(cmd)
276 proc.SetProcessGroupKill(cmd)
277 cmd.Cancel = func() error {
278 proc.KillTree(cmd)
279 return nil
280 }
281 cmd.WaitDelay = probeWaitDelay
282 }
283
284 func sortResults(results []ProbeResult) {
285 sort.Slice(results, func(i, j int) bool {
286 if results[i].Found != results[j].Found {
287 return results[i].Found
288 }
289 return results[i].Binary < results[j].Binary
290 })
291 }
292
293 func firstLine(s string) string {
294 s = strings.TrimSpace(s)
295 if i := strings.IndexByte(s, '\n'); i >= 0 {
296 return strings.TrimRight(s[:i], "\r")
297 }
298 return s
299 }
300
301 func FormatSection(results []ProbeResult, osName, shellPath string, overrides map[string]string) string {
302 if len(results) == 0 && len(overrides) == 0 && osName == "" && shellPath == "" {
303 return ""
304 }
305 results = append([]ProbeResult(nil), results...)
306 sortResults(results)
307 var b strings.Builder
308 b.WriteString("## Environment\n\n")
309 if osName == "" {
310 osName = runtime.GOOS + "/" + runtime.GOARCH
311 }
312 b.WriteString("- OS: " + osName + "\n")
313 if shellPath != "" {
314 b.WriteString("- Shell: " + redactHome(shellPath) + "\n")
315 }
316 if len(overrides) > 0 {
317 b.WriteString("\nConfigured tools:\n")
318 names := sortedMapKeys(overrides)
319 for _, name := range limitStrings(names, maxRenderedTools) {
320 fmt.Fprintf(&b, "- %s: %s\n", name, redactHome(overrides[name]))
321 }
322 if omitted := len(names) - maxRenderedTools; omitted > 0 {
323 fmt.Fprintf(&b, "- ... %d more configured tools omitted\n", omitted)
324 }
325 }
326 if len(results) > 0 {
327 b.WriteString("\nDetected tools:\n")
328 foundShown := 0
329 foundTotal := 0
330 for _, r := range results {
331 if r.Found {
332 foundTotal++
333 if foundShown >= maxRenderedTools {
334 continue
335 }
336 out := r.Output
337 if out == "" {
338 out = "available"
339 }
340 fmt.Fprintf(&b, "- %s: %s\n", r.Binary, out)
341 foundShown++
342 }
343 }
344 if omitted := foundTotal - foundShown; omitted > 0 {
345 fmt.Fprintf(&b, "- ... %d more detected tools omitted\n", omitted)
346 }
347 b.WriteString("\nNot found or unavailable:\n")
348 missingShown := 0
349 missingTotal := 0
350 for _, r := range results {
351 if !r.Found {
352 missingTotal++
353 if missingShown >= maxRenderedTools {
354 continue
355 }
356 reason := r.Error
357 if reason == "" {
358 reason = "not found"
359 }
360 fmt.Fprintf(&b, "- %s: %s\n", r.Binary, reason)
361 missingShown++
362 }
363 }
364 if omitted := missingTotal - missingShown; omitted > 0 {
365 fmt.Fprintf(&b, "- ... %d more unavailable tools omitted\n", omitted)
366 }
367 }
368 b.WriteString("\nUse detected tools when appropriate. Do not try unavailable tools unless the user installs or configures them.\n")
369 return strings.TrimRight(b.String(), "\n")
370 }
371
372 func sortedMapKeys(m map[string]string) []string {
373 keys := make([]string, 0, len(m))
374 for k := range m {
375 if strings.TrimSpace(k) != "" {
376 keys = append(keys, k)
377 }
378 }
379 sort.Strings(keys)
380 return keys
381 }
382
383 func limitStrings(in []string, limit int) []string {
384 if len(in) <= limit {
385 return in
386 }
387 return in[:limit]
388 }
389
390 func redactHome(path string) string {
391 path = expandHome(path)
392 if path == "" {
393 return ""
394 }
395 home, err := os.UserHomeDir()
396 if err != nil || home == "" {
397 return filepath.ToSlash(filepath.Clean(path))
398 }
399 clean := filepath.Clean(path)
400 home = filepath.Clean(home)
401 if clean == home {
402 return "~"
403 }
404 if strings.HasPrefix(clean, home+string(filepath.Separator)) {
405 return filepath.ToSlash("~" + strings.TrimPrefix(clean, home))
406 }
407 return filepath.ToSlash(clean)
408 }
409
410 func expandHome(path string) string {
411 path = strings.TrimSpace(path)
412 if path == "" || path == "~" || !strings.HasPrefix(path, "~/") {
413 return path
414 }
415 home, err := os.UserHomeDir()
416 if err != nil || home == "" {
417 return path
418 }
419 return filepath.Join(home, strings.TrimPrefix(path, "~/"))
420 }
421
422 func fileExecutable(path string) bool {
423 fi, err := os.Stat(path)
424 return err == nil && !fi.IsDir()
425 }
426
427 func blockedExecutable(path string, denyRoots []string) bool {
428 abs, err := filepath.Abs(path)
429 if err != nil {
430 abs = filepath.Clean(path)
431 }
432 for _, root := range normalizedDenyRoots(denyRoots) {
433 if pathWithin(abs, root) {
434 return true
435 }
436 }
437 return false
438 }
439
440 func normalizedDenyRoots(roots []string) []string {
441 out := make([]string, 0, len(roots))
442 seen := map[string]bool{}
443 for _, root := range roots {
444 root = strings.TrimSpace(root)
445 if root == "" {
446 continue
447 }
448 root = expandHome(root)
449 abs, err := filepath.Abs(root)
450 if err != nil {
451 abs = filepath.Clean(root)
452 }
453 abs = filepath.Clean(abs)
454 if !seen[abs] {
455 seen[abs] = true
456 out = append(out, abs)
457 }
458 }
459 sort.Strings(out)
460 return out
461 }
462
463 func pathWithin(path, root string) bool {
464 path = filepath.Clean(path)
465 root = filepath.Clean(root)
466 if path == root {
467 return true
468 }
469 rel, err := filepath.Rel(root, path)
470 if err != nil {
471 return false
472 }
473 return rel != "." && rel != "" && !strings.HasPrefix(rel, ".."+string(filepath.Separator)) && rel != ".."
474 }
475
475 lines GO