diff --git a/cmd/run.go b/cmd/run.go index 7b2b20e..5077a72 100644 --- a/cmd/run.go +++ b/cmd/run.go @@ -68,13 +68,13 @@ func runHandler(cmd *cobra.Command, args []string) error { dbPath, err := config.DBPath() if err != nil { logFn("[warn] cannot open cache: %v", err) - return runWithoutCache(cfg, proj, wm, logFn) + return runWithoutCache(cfg, proj, wm, postCompact, logFn) } store, err := cache.Open(dbPath) if err != nil { logFn("[warn] cache open error: %v", err) - return runWithoutCache(cfg, proj, wm, logFn) + return runWithoutCache(cfg, proj, wm, postCompact, logFn) } defer store.Close() @@ -192,7 +192,7 @@ func runHandler(cmd *cobra.Command, args []string) error { } // runWithoutCache attempts an API fetch with no cache fallback. -func runWithoutCache(cfg *config.Config, proj *project.Info, wm *project.WorkingMemory, logFn func(string, ...interface{})) error { +func runWithoutCache(cfg *config.Config, proj *project.Info, wm *project.WorkingMemory, postCompact bool, logFn func(string, ...interface{})) error { ctx, cancel := context.WithTimeout(context.Background(), 90*time.Second) defer cancel() @@ -218,7 +218,7 @@ func runWithoutCache(cfg *config.Config, proj *project.Info, wm *project.Working return silentExit() } - opts := tmpl.RenderOptions{MaxTokens: maxTokens, WorkingMemory: wm} + opts := tmpl.RenderOptions{MaxTokens: maxTokens, WorkingMemory: wm, PostCompact: postCompact} output, _, err := tmpl.Render(graph, proj.Name, opts) if err != nil { logFn("[warn] render error: %v", err) diff --git a/internal/hooks/hooks.go b/internal/hooks/hooks.go index 1b121f4..582a04a 100644 --- a/internal/hooks/hooks.go +++ b/internal/hooks/hooks.go @@ -187,7 +187,7 @@ func commandExistsInHooks(hookList []Hook, matches ...string) bool { // isAlreadyInstalled checks if ALL uncompact hooks are present. func isAlreadyInstalled(hooks map[string][]Hook) bool { return commandExistsInHooks(hooks["Stop"], "uncompact run", "uncompact-hook.sh") && - commandExistsInHooks(hooks["UserPromptSubmit"], "uncompact show-cache") + commandExistsInHooks(hooks["UserPromptSubmit"], "uncompact show-cache", "show-hook.sh") } // mergeHooks adds hooks from toAdd into existing, skipping any whose commands diff --git a/internal/template/render.go b/internal/template/render.go index 87e5e67..3fb7e46 100644 --- a/internal/template/render.go +++ b/internal/template/render.go @@ -137,7 +137,27 @@ func Render(graph *api.ProjectGraph, projectName string, opts RenderOptions) (st } if opts.PostCompact { - note := fmt.Sprintf("\n\n> [uncompact] Context restored after compact (~%d tokens). Please acknowledge at the start of your response with exactly: \"✓ Uncompact: context restored (~%d tokens)\"", resultTokens, resultTokens) + buildNote := func(n int) string { + return fmt.Sprintf("\n\n> [uncompact] Context restored after compact (~%d tokens). Please acknowledge at the start of your response with exactly: \"✓ Uncompact: context restored (~%d tokens)\"", n, n) + } + note := buildNote(resultTokens) + noteTokens := countTokens(note) + if opts.MaxTokens > 0 && resultTokens+noteTokens > opts.MaxTokens { + budget := opts.MaxTokens - noteTokens + if budget < 1 { + budget = 1 + } + if tokens <= budget { + result, resultTokens = fullText, tokens + } else { + truncated, truncatedTokens, truncErr := truncateToTokenBudget(graph, projectName, budget, graph.Stats.CircularDependencyCycles, opts.WorkingMemory) + if truncErr != nil { + return "", 0, truncErr + } + result, resultTokens = truncated, truncatedTokens + } + note = buildNote(resultTokens) + } result += note resultTokens = countTokens(result) } diff --git a/scripts/uncompact-hook.sh b/scripts/uncompact-hook.sh index cd1885f..f73445d 100644 --- a/scripts/uncompact-hook.sh +++ b/scripts/uncompact-hook.sh @@ -38,10 +38,9 @@ OUTPUT="$("$UNCOMPACT" run --fallback --post-compact)" DISPLAY_CACHE="${TMPDIR:-/tmp}/uncompact-display-${UID:-$(id -u)}.txt" if [ -n "$OUTPUT" ]; then - CHAR_COUNT="${#OUTPUT}" - APPROX_TOKENS=$(( CHAR_COUNT / 4 )) - - # Write to display cache — UserPromptSubmit hook will show this as a visible - # transcript message on the user's next message. - printf '%s\n\n[uncompact] Context restored (~%d tokens)\n' "$OUTPUT" "$APPROX_TOKENS" > "$DISPLAY_CACHE" + # Write securely and atomically to avoid disclosure/races. + umask 077 + TMP_CACHE="$(mktemp "${DISPLAY_CACHE}.XXXXXX")" || exit 0 + printf '%s\n' "$OUTPUT" > "$TMP_CACHE" + mv -f "$TMP_CACHE" "$DISPLAY_CACHE" fi