diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 8eb2e52..5e9a18e 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -11,7 +11,7 @@ concurrency: cancel-in-progress: true env: - GO_VERSION: '1.23' + GO_VERSION: '1.25' jobs: test: @@ -69,7 +69,11 @@ jobs: - name: Run tests run: | - go test -v -race -coverprofile=coverage.out -covermode=atomic ./pkg/... + # The whole module: ./pkg/... alone skips the LangGraph conformance + # suite and the Studio end-to-end suite under ./test/.... + go test -v -race -timeout 15m \ + -coverpkg=./pkg/... -coverprofile=coverage.out -covermode=atomic \ + ./... env: POSTGRES_HOST: localhost POSTGRES_PORT: 5432 @@ -79,6 +83,28 @@ jobs: REDIS_HOST: localhost REDIS_PORT: 6379 + - name: Fuzz smoke test + run: | + # A short run of each target: enough to catch a regression that makes a + # fuzz target crash immediately, without lengthening CI. + set -e + for pkg_target in \ + "./pkg/core FuzzStateJSONRoundTrip" \ + "./pkg/core FuzzStateMarshalUnmarshal" \ + "./pkg/core FuzzDeepCopy" \ + "./pkg/core FuzzGraphRouting" \ + "./pkg/core FuzzGraphConstruction" \ + "./pkg/tools FuzzToolArguments" \ + "./pkg/tools FuzzPathResolution" \ + "./pkg/tools FuzzCommandPolicy" \ + "./pkg/tools FuzzURLPolicy" \ + "./pkg/server FuzzAPIRequestBodies" \ + "./pkg/server FuzzAPIPaths" ; do + set -- $pkg_target + echo "fuzzing $2 in $1" + go test -run '^$' -fuzz "^$2$" -fuzztime=20s "$1" + done + - name: Upload coverage to Codecov uses: codecov/codecov-action@v4 with: @@ -245,7 +271,7 @@ jobs: - name: Run integration tests run: | - go test -v -tags=integration ./test/e2e/... + go test -v -race -timeout 15m ./test/e2e/... env: POSTGRES_HOST: localhost POSTGRES_PORT: 5432 diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml index 3988667..c6c4227 100644 --- a/.github/workflows/docs.yml +++ b/.github/workflows/docs.yml @@ -24,7 +24,7 @@ concurrency: cancel-in-progress: true env: - GO_VERSION: '1.23' + GO_VERSION: '1.25' jobs: build: diff --git a/.github/workflows/pre-commit.yml b/.github/workflows/pre-commit.yml index ee1f1b5..1102731 100644 --- a/.github/workflows/pre-commit.yml +++ b/.github/workflows/pre-commit.yml @@ -12,7 +12,7 @@ concurrency: cancel-in-progress: true env: - GO_VERSION: '1.23' + GO_VERSION: '1.25' jobs: pre-commit: @@ -118,7 +118,7 @@ jobs: continue-on-error: true - name: Run Trivy vulnerability scanner - uses: aquasecurity/trivy-action@0.25.0 + uses: aquasecurity/trivy-action@v0.36.0 with: scan-type: 'fs' scan-ref: '.' diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 8244da4..8e9e4fd 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -23,7 +23,7 @@ concurrency: cancel-in-progress: true env: - GO_VERSION: '1.23' + GO_VERSION: '1.25' REGISTRY: ghcr.io IMAGE_NAME: ${{ github.repository }} diff --git a/.gitignore b/.gitignore index a122ca1..a746a0a 100644 --- a/.gitignore +++ b/.gitignore @@ -6,6 +6,18 @@ bin/ *.so *.dylib +# Compiled example binaries. These were previously committed, adding ~100 MB of +# platform-specific artifacts to every clone that could silently drift from the +# source they were built from. Build them with `go build ./...` in the example. +/examples/*/[0-9][0-9]-* +!/examples/*/[0-9][0-9]-*/ +/examples/01-basic-chat/basic-chat +/examples/02-react-agent/react-agent +/examples/07-tools-integration/tools-integration +/examples/08-production-ready/production-ready +/examples/09-workflow-graph/workflow-graph +/examples/streaming-demo/streaming-demo + # Test binary, built with `go test -c` *.test diff --git a/Dockerfile b/Dockerfile index 7cc57bf..8f6ca63 100644 --- a/Dockerfile +++ b/Dockerfile @@ -1,5 +1,5 @@ # Build stage -FROM golang:1.23.10-alpine AS builder +FROM golang:1.25.13-alpine AS builder # Set working directory WORKDIR /app @@ -47,9 +47,14 @@ USER golanggraph # Expose port (adjust as needed) EXPOSE 8080 -# Health check -HEALTHCHECK --interval=30s --timeout=3s --start-period=5s --retries=3 \ - CMD ./golanggraph health || exit 1 +# Health check. +# +# Probes the server's own endpoint rather than running a local dependency +# scan: the container is healthy when it is serving. A plain "health" run +# reports on dependencies, which is a different question and would mark a +# perfectly serving container unhealthy whenever an optional one is absent. +HEALTHCHECK --interval=30s --timeout=5s --start-period=10s --retries=3 \ + CMD ["./golanggraph", "health", "--server", "http://127.0.0.1:8080"] # Run the binary ENTRYPOINT ["./golanggraph"] diff --git a/Dockerfile.agent b/Dockerfile.agent index 8e9f594..17bc47b 100644 --- a/Dockerfile.agent +++ b/Dockerfile.agent @@ -1,5 +1,5 @@ # Production Dockerfile for GoLangGraph Agent -FROM golang:1.23.10-alpine AS builder +FROM golang:1.25.13-alpine AS builder # Set working directory WORKDIR /app @@ -37,12 +37,10 @@ WORKDIR /app # Copy the binary from builder stage COPY --from=builder /app/golanggraph-agent . -# Copy configuration files -COPY configs/ ./configs/ -COPY static/ ./static/ - -# Create necessary directories and change ownership to non-root user -RUN mkdir -p ./logs ./data && \ +# Create the directories the server may use. These were previously COPY'd from +# configs/ and static/, which do not exist in the repository, so this image +# could not be built at all. Mount or bake in real configuration as needed. +RUN mkdir -p ./configs ./static ./logs ./data && \ chown -R golanggraph:golanggraph /app # Switch to non-root user @@ -51,9 +49,13 @@ USER golanggraph # Expose port EXPOSE 8080 -# Health check -HEALTHCHECK --interval=30s --timeout=10s --start-period=5s --retries=3 \ - CMD ./golanggraph-agent health || exit 1 +# Health check. +# +# Probes the server's own endpoint: the container is healthy when it is +# serving. A local dependency scan is a different question and would mark a +# perfectly serving container unhealthy whenever an optional service is absent. +HEALTHCHECK --interval=30s --timeout=10s --start-period=10s --retries=3 \ + CMD ["./golanggraph-agent", "health", "--server", "http://127.0.0.1:8080"] # Run the agent ENTRYPOINT ["./golanggraph-agent"] diff --git a/README.md b/README.md index d717c1f..2527e4c 100644 --- a/README.md +++ b/README.md @@ -173,6 +173,26 @@ if err != nil { fmt.Printf("πŸ”„ Graph Result: %v\n", result.Get("response")) ``` +## πŸ”’ Production Deployment + +Defaults favour local development. Before exposing GoLangGraph to real traffic, +read **[docs/PRODUCTION.md](docs/PRODUCTION.md)**, which covers: + +- **Authentication and CORS** β€” `RequireAuth` is off by default, and the allowed-origin list also governs WebSocket upgrades. +- **Tool sandboxing** β€” filesystem confinement, the shell allowlist, and SSRF protection for the HTTP tool. +- **Durable execution** β€” checkpointing, resume after a crash, and human-in-the-loop interrupts. +- **Health checking** β€” which probe belongs in a container, and which does not. +- **Typed errors, retries, concurrency and observability.** + +LangGraph compatibility, including the places GoLangGraph intentionally +differs, is documented in +**[test/conformance/DEVIATIONS.md](test/conformance/DEVIATIONS.md)** and +enforced by the conformance suite: + +```bash +go test -race ./test/conformance/... +``` + ## πŸ—οΈ Architecture GoLangGraph follows a modular architecture: diff --git a/cmd/golanggraph/auto_serve_command.go b/cmd/golanggraph/auto_serve_command.go index 76435b9..23dc51a 100644 --- a/cmd/golanggraph/auto_serve_command.go +++ b/cmd/golanggraph/auto_serve_command.go @@ -7,12 +7,15 @@ package main import ( - "context" "fmt" + "io" "os" "os/signal" "path/filepath" + "sort" + "strings" "syscall" + "time" "github.com/spf13/cobra" @@ -80,14 +83,13 @@ func init() { // Development features autoServeCmd.Flags().Bool("dev", false, "Enable development mode") - autoServeCmd.Flags().Bool("watch", false, "Watch for file changes and hot-reload") - autoServeCmd.Flags().Bool("verbose", false, "Enable verbose logging") + autoServeCmd.Flags().Bool("watch", false, "Watch for file changes and hot-reload (not implemented)") // Production features autoServeCmd.Flags().String("env", "development", "Environment (development, staging, production)") - autoServeCmd.Flags().String("log-level", "info", "Log level (debug, info, warn, error)") - autoServeCmd.Flags().Duration("timeout", 0, "Request timeout") - autoServeCmd.Flags().Int64("max-request-size", 0, "Maximum request size in bytes") + autoServeCmd.Flags().String("log-level", "info", "Log level (not applied: the auto-server logs at info)") + autoServeCmd.Flags().Duration("timeout", 30*time.Second, "Request read/write timeout") + autoServeCmd.Flags().Int64("max-request-size", 10*1024*1024, "Maximum request size in bytes") // Docker and deployment autoServeCmd.Flags().Bool("generate-dockerfile", false, "Generate Dockerfile for deployment") @@ -99,248 +101,372 @@ func init() { autoServeCmd.Flags().StringSlice("agent-dirs", []string{}, "Additional agent directories") } +// autoServeOptions is the resolved configuration of an auto-serve run. +type autoServeOptions struct { + SourcePath string + Env string + Dev bool + Watch bool + LogLevel string + AgentDirs []string + Plugins []string + GenerateDockerfile bool + GenerateDockerCompose bool + GenerateK8s bool +} + func runAutoServe(cmd *cobra.Command, args []string) error { - // Parse flags - host, _ := cmd.Flags().GetString("host") - port, _ := cmd.Flags().GetInt("port") - basePath, _ := cmd.Flags().GetString("base-path") - webUI, _ := cmd.Flags().GetBool("web-ui") - playground, _ := cmd.Flags().GetBool("playground") - schemaAPI, _ := cmd.Flags().GetBool("schema-api") - metrics, _ := cmd.Flags().GetBool("metrics") - cors, _ := cmd.Flags().GetBool("cors") - schemaValidation, _ := cmd.Flags().GetBool("schema-validation") - ollamaEndpoint, _ := cmd.Flags().GetString("ollama-endpoint") - dev, _ := cmd.Flags().GetBool("dev") - watch, _ := cmd.Flags().GetBool("watch") - _, _ = cmd.Flags().GetBool("verbose") // verbose flag parsed but not used in current implementation - env, _ := cmd.Flags().GetString("env") - generateDockerfile, _ := cmd.Flags().GetBool("generate-dockerfile") - generateDockerCompose, _ := cmd.Flags().GetBool("generate-docker-compose") - generateK8s, _ := cmd.Flags().GetBool("generate-k8s") - plugins, _ := cmd.Flags().GetStringSlice("plugins") - agentDirs, _ := cmd.Flags().GetStringSlice("agent-dirs") - - fmt.Printf("πŸš€ GoLangGraph Auto-Serve - Minimal Code, Maximum Power!\n") - fmt.Printf("═══════════════════════════════════════════════════════════\n\n") - - // Determine source path - sourcePath := "." - if len(args) > 0 { - sourcePath = args[0] + out := cmd.OutOrStdout() + + config, opts, err := autoServeConfigFromFlags(cmd, args) + if err != nil { + return err + } + + autoServer, err := prepareAutoServer(out, config, opts) + if err != nil { + return err } - // Create auto-server configuration + // Fail before announcing a listening server if the address is unusable. + if err := checkAddressAvailable(config.Host, config.Port); err != nil { + return err + } + + ctx, stop := signal.NotifyContext(cmd.Context(), syscall.SIGINT, syscall.SIGTERM) + defer stop() + + printAutoServeURLs(out, config) + + if err := autoServer.Start(ctx); err != nil { + return fmt.Errorf("server failed: %w", err) + } + + _, _ = fmt.Fprintf(out, "βœ… Server stopped gracefully\n") + return nil +} + +// autoServeConfigFromFlags turns the command's flags into a server +// configuration and run options. +// +// Keeping this out of the serving loop makes it testable that every declared +// flag reaches the configuration: --timeout and --max-request-size were +// declared and then never read, so the server ran with neither applied. +func autoServeConfigFromFlags(cmd *cobra.Command, args []string) (*server.AutoServerConfig, autoServeOptions, error) { config := &server.AutoServerConfig{ - Host: host, - Port: port, - BasePath: basePath, - EnableWebUI: webUI, - EnablePlayground: playground, - EnableSchemaAPI: schemaAPI, - EnableMetricsAPI: metrics, - EnableCORS: cors, - SchemaValidation: schemaValidation, - OllamaEndpoint: ollamaEndpoint, - LLMProviders: make(map[string]interface{}), - Middleware: []string{"cors", "logging", "recovery"}, - } - - // Add LLM providers based on flags - if openaiKey, _ := cmd.Flags().GetString("openai-api-key"); openaiKey != "" { + LLMProviders: make(map[string]interface{}), + Middleware: []string{"cors", "logging", "recovery"}, + } + opts := autoServeOptions{SourcePath: "."} + if len(args) > 0 { + opts.SourcePath = args[0] + } + + // Flag reads are checked: a mistyped flag name used to be swallowed by + // `value, _ := cmd.Flags().Get...` and silently produce a zero value. + flags := cmd.Flags() + var err error + if config.Host, err = flags.GetString("host"); err != nil { + return nil, opts, err + } + if config.Port, err = flags.GetInt("port"); err != nil { + return nil, opts, err + } + if config.BasePath, err = flags.GetString("base-path"); err != nil { + return nil, opts, err + } + if config.EnableWebUI, err = flags.GetBool("web-ui"); err != nil { + return nil, opts, err + } + if config.EnablePlayground, err = flags.GetBool("playground"); err != nil { + return nil, opts, err + } + if config.EnableSchemaAPI, err = flags.GetBool("schema-api"); err != nil { + return nil, opts, err + } + if config.EnableMetricsAPI, err = flags.GetBool("metrics"); err != nil { + return nil, opts, err + } + if config.EnableCORS, err = flags.GetBool("cors"); err != nil { + return nil, opts, err + } + if config.SchemaValidation, err = flags.GetBool("schema-validation"); err != nil { + return nil, opts, err + } + if config.OllamaEndpoint, err = flags.GetString("ollama-endpoint"); err != nil { + return nil, opts, err + } + if config.ServerTimeout, err = flags.GetDuration("timeout"); err != nil { + return nil, opts, err + } + if config.MaxRequestSize, err = flags.GetInt64("max-request-size"); err != nil { + return nil, opts, err + } + + openaiKey, err := flags.GetString("openai-api-key") + if err != nil { + return nil, opts, err + } + if openaiKey != "" { config.LLMProviders["openai"] = map[string]string{"api_key": openaiKey} } - if anthropicKey, _ := cmd.Flags().GetString("anthropic-api-key"); anthropicKey != "" { + anthropicKey, err := flags.GetString("anthropic-api-key") + if err != nil { + return nil, opts, err + } + if anthropicKey != "" { config.LLMProviders["anthropic"] = map[string]string{"api_key": anthropicKey} } - // Adjust config for environment - if env == "production" { - config.EnablePlayground = false // Disable playground in production - fmt.Printf("πŸ”’ Production mode enabled - some debug features disabled\n") + if opts.Env, err = flags.GetString("env"); err != nil { + return nil, opts, err + } + if opts.Dev, err = flags.GetBool("dev"); err != nil { + return nil, opts, err + } + if opts.Watch, err = flags.GetBool("watch"); err != nil { + return nil, opts, err + } + if opts.LogLevel, err = flags.GetString("log-level"); err != nil { + return nil, opts, err + } + if opts.AgentDirs, err = flags.GetStringSlice("agent-dirs"); err != nil { + return nil, opts, err + } + if opts.Plugins, err = flags.GetStringSlice("plugins"); err != nil { + return nil, opts, err + } + if opts.GenerateDockerfile, err = flags.GetBool("generate-dockerfile"); err != nil { + return nil, opts, err + } + if opts.GenerateDockerCompose, err = flags.GetBool("generate-docker-compose"); err != nil { + return nil, opts, err + } + if opts.GenerateK8s, err = flags.GetBool("generate-k8s"); err != nil { + return nil, opts, err } - if dev { - fmt.Printf("πŸ› οΈ Development mode enabled\n") - if watch { - fmt.Printf("πŸ‘€ File watching enabled (hot-reload)\n") - } + return config, opts, nil +} + +// prepareAutoServer resolves the configuration, loads the agents and generates +// any deployment files, without binding a port. Keeping this separate from the +// serving loop is what makes the command testable. +func prepareAutoServer(out io.Writer, config *server.AutoServerConfig, opts autoServeOptions) (*server.AutoServer, error) { + _, _ = fmt.Fprintf(out, "πŸš€ GoLangGraph Auto-Serve\n") + + if config.Port <= 0 || config.Port > 65535 { + return nil, fmt.Errorf("invalid port %d", config.Port) + } + switch opts.Env { + case "development", "staging", "production": + default: + return nil, fmt.Errorf("unknown environment %q (want development, staging or production)", opts.Env) } - // Create auto-server - autoServer := server.NewAutoServer(config) + if opts.Env == "production" { + config.EnablePlayground = false + _, _ = fmt.Fprintf(out, "πŸ”’ Production mode: playground disabled\n") + } + if opts.Dev { + _, _ = fmt.Fprintf(out, "πŸ› οΈ Development mode enabled\n") + } + if opts.Watch { + // This used to print "πŸ‘€ File watching enabled (hot-reload)". Nothing + // ever watched anything. + _, _ = fmt.Fprintf(out, "⚠️ --watch is not implemented; restart the server to pick up changes\n") + } + if opts.LogLevel != "" && opts.LogLevel != "info" { + _, _ = fmt.Fprintf(out, "⚠️ --log-level is not applied; the auto-server logs at info level\n") + } - // Load agents from various sources - fmt.Printf("πŸ“ Loading agents from: %s\n", sourcePath) + autoServer := server.NewAutoServer(config) + registry := agent.GetGlobalRegistry() + before := len(registry.ListDefinitions()) + + // A source path that does not exist used to be skipped in silence: the + // command then served three example agents while reporting that it had + // loaded agents from the operator's path. + sources := append([]string{opts.SourcePath}, opts.AgentDirs...) + for i, source := range sources { + info, err := os.Stat(source) + if err != nil { + if i == 0 && opts.SourcePath == "." { + // No explicit path given and no current directory: nothing to load. + continue + } + return nil, fmt.Errorf("agent source %s: %w", source, err) + } - // Check if source is a file or directory - if info, err := os.Stat(sourcePath); err == nil { if info.IsDir() { - // Load from directory - if err := autoServer.LoadAgentsFromDirectory(sourcePath); err != nil { - fmt.Printf("⚠️ Warning: %v\n", err) - } - } else if filepath.Ext(sourcePath) == ".yaml" || filepath.Ext(sourcePath) == ".yml" { - // Load from config file - if err := autoServer.LoadAgentsFromConfig(sourcePath); err != nil { - return fmt.Errorf("failed to load agents from config: %w", err) + loaded, err := loadAgentsFromDirectory(out, autoServer, source) + if err != nil { + return nil, err } + _, _ = fmt.Fprintf(out, "πŸ“ %s: %d agent config file(s) loaded\n", source, loaded) + continue } - } - // Load additional agent directories - for _, dir := range agentDirs { - fmt.Printf("πŸ“ Loading additional agents from: %s\n", dir) - if err := autoServer.LoadAgentsFromDirectory(dir); err != nil { - fmt.Printf("⚠️ Warning: %v\n", err) + switch ext := strings.ToLower(filepath.Ext(source)); ext { + case ".yaml", ".yml", ".json": + if err := autoServer.LoadAgentsFromConfig(source); err != nil { + return nil, fmt.Errorf("failed to load agents from config: %w", err) + } + _, _ = fmt.Fprintf(out, "πŸ“„ %s: loaded\n", source) + default: + // Previously any other extension was ignored without a word. + return nil, fmt.Errorf("unsupported agent source %s: want a directory or a .yaml, .yml or .json file", source) } } - // Load plugins - registry := agent.GetGlobalRegistry() - for _, pluginPath := range plugins { - fmt.Printf("πŸ”Œ Loading plugin: %s\n", pluginPath) + for _, pluginPath := range opts.Plugins { + _, _ = fmt.Fprintf(out, "πŸ”Œ Loading plugin: %s\n", pluginPath) if err := registry.LoadFromPlugin(pluginPath); err != nil { - fmt.Printf("⚠️ Warning: Failed to load plugin %s: %v\n", pluginPath, err) + return nil, fmt.Errorf("failed to load plugin %s: %w", pluginPath, err) } } - // Register example agents if none found - if len(registry.ListDefinitions()) == 0 { - fmt.Printf("πŸ“ No agents found, creating example agents...\n") - createExampleAgents(autoServer) + definitions := registry.ListDefinitions() + if len(definitions) == 0 { + _, _ = fmt.Fprintf(out, "πŸ“ No agents found; registering the built-in example agents\n") + if err := createExampleAgents(out, autoServer); err != nil { + return nil, err + } + } else { + _, _ = fmt.Fprintf(out, "πŸ€– %d agent(s) registered (%d from this run)\n", + len(definitions), len(definitions)-before) } - // Generate deployment files if requested - if generateDockerfile { - if err := generateDockerfileForProject(sourcePath); err != nil { - fmt.Printf("⚠️ Warning: Failed to generate Dockerfile: %v\n", err) - } else { - fmt.Printf("🐳 Generated Dockerfile\n") + if opts.GenerateDockerfile { + if err := generateDockerfileForProject(opts.SourcePath); err != nil { + return nil, fmt.Errorf("failed to generate Dockerfile: %w", err) } + _, _ = fmt.Fprintf(out, "🐳 Generated Dockerfile\n") } - - if generateDockerCompose { - if err := generateDockerComposeForProject(sourcePath, config); err != nil { - fmt.Printf("⚠️ Warning: Failed to generate docker-compose.yml: %v\n", err) - } else { - fmt.Printf("🐳 Generated docker-compose.yml\n") + if opts.GenerateDockerCompose { + if err := generateDockerComposeForProject(opts.SourcePath, config); err != nil { + return nil, fmt.Errorf("failed to generate docker-compose.yml: %w", err) } + _, _ = fmt.Fprintf(out, "🐳 Generated docker-compose.yml\n") } - - if generateK8s { - if err := generateKubernetesManifests(sourcePath, config); err != nil { - fmt.Printf("⚠️ Warning: Failed to generate Kubernetes manifests: %v\n", err) - } else { - fmt.Printf("☸️ Generated Kubernetes manifests\n") + if opts.GenerateK8s { + if err := generateKubernetesManifests(opts.SourcePath, config); err != nil { + return nil, fmt.Errorf("failed to generate Kubernetes manifests: %w", err) } + _, _ = fmt.Fprintf(out, "☸️ Generated Kubernetes manifests\n") } - // Print configuration summary - fmt.Printf("\nπŸ”§ Configuration:\n") - fmt.Printf(" Host: %s\n", host) - fmt.Printf(" Port: %d\n", port) - fmt.Printf(" Base Path: %s\n", basePath) - fmt.Printf(" Environment: %s\n", env) - fmt.Printf(" Ollama: %s\n", ollamaEndpoint) - fmt.Printf(" Features: ") - features := []string{} - if webUI { - features = append(features, "WebUI") - } - if playground { - features = append(features, "Playground") - } - if schemaAPI { - features = append(features, "Schema API") + return autoServer, nil +} + +// loadAgentsFromDirectory registers the agents defined by the configuration +// files in a directory. +// +// server.AutoServer.LoadAgentsFromDirectory is a stub that loads nothing and +// returns nil, so relying on it meant "πŸ“ Loading agents from: ./agents" +// followed by an empty server. The configuration files are loaded here instead. +func loadAgentsFromDirectory(out io.Writer, autoServer *server.AutoServer, dir string) (int, error) { + entries, err := os.ReadDir(dir) + if err != nil { + return 0, err } - if metrics { - features = append(features, "Metrics") + + names := make([]string, 0, len(entries)) + for _, entry := range entries { + if entry.IsDir() { + continue + } + switch strings.ToLower(filepath.Ext(entry.Name())) { + case ".yaml", ".yml", ".json": + names = append(names, entry.Name()) + } } - if schemaValidation { - features = append(features, "Validation") + sort.Strings(names) + + loaded := 0 + for _, name := range names { + path := filepath.Join(dir, name) + if err := autoServer.LoadAgentsFromConfig(path); err != nil { + // Not every YAML file in a directory is an agent config; report it + // and carry on rather than refusing to start. + _, _ = fmt.Fprintf(out, " ⚠️ %s: %v\n", path, err) + continue + } + loaded++ } - fmt.Printf("%v\n\n", features) - - // Setup graceful shutdown - ctx, cancel := context.WithCancel(context.Background()) - - // Handle shutdown signals - sigChan := make(chan os.Signal, 1) - signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM) - go func() { - <-sigChan - fmt.Printf("\nπŸ›‘ Shutdown signal received, stopping server...\n") - cancel() - }() + return loaded, nil +} - // Start the server - fmt.Printf("πŸŽ‰ Starting auto-generated multi-agent system!\n") - fmt.Printf("═══════════════════════════════════════════════════════════\n") +// printAutoServeURLs prints where to reach the server. +func printAutoServeURLs(out io.Writer, config *server.AutoServerConfig) { + host := config.Host + if host == "0.0.0.0" || host == "" { + host = "localhost" + } + baseURL := fmt.Sprintf("http://%s:%d", host, config.Port) - // Print quick access URLs - baseURL := fmt.Sprintf("http://localhost:%d", port) - fmt.Printf("\n🌐 Quick Access URLs:\n") - if webUI { - fmt.Printf(" πŸ’¬ Chat Interface: %s/chat\n", baseURL) + _, _ = fmt.Fprintf(out, "\n🌐 URLs:\n") + if config.EnableWebUI { + _, _ = fmt.Fprintf(out, " πŸ’¬ Chat Interface: %s/chat\n", baseURL) } - if playground { - fmt.Printf(" πŸ—οΈ API Playground: %s/playground\n", baseURL) + if config.EnablePlayground { + _, _ = fmt.Fprintf(out, " πŸ—οΈ API Playground: %s/playground\n", baseURL) } - fmt.Printf(" πŸ“‹ System Health: %s/health\n", baseURL) - fmt.Printf(" πŸ€– List Agents: %s/agents\n", baseURL) - if schemaAPI { - fmt.Printf(" πŸ“„ API Schemas: %s/schemas\n", baseURL) + _, _ = fmt.Fprintf(out, " πŸ“‹ System Health: %s/health\n", baseURL) + _, _ = fmt.Fprintf(out, " πŸ€– List Agents: %s/agents\n", baseURL) + if config.EnableSchemaAPI { + _, _ = fmt.Fprintf(out, " πŸ“„ API Schemas: %s/schemas\n", baseURL) } - if metrics { - fmt.Printf(" πŸ“Š Metrics: %s/metrics\n", baseURL) + if config.EnableMetricsAPI { + _, _ = fmt.Fprintf(out, " πŸ“Š Metrics: %s/metrics\n", baseURL) } - fmt.Printf(" πŸ”§ Debug: %s/debug\n", baseURL) - fmt.Printf("\n") + _, _ = fmt.Fprintf(out, "\n") +} - if err := autoServer.Start(ctx); err != nil { - return fmt.Errorf("server failed: %w", err) +// defaultExampleModel is the model the example agents are configured with; it +// is small enough to run under a local Ollama. +const defaultExampleModel = "gemma3:1b" + +// createExampleAgents registers the demonstration agents used when no agent +// definitions were found. +func createExampleAgents(out io.Writer, autoServer *server.AutoServer) error { + examples := []struct { + id, name, prompt string + agentType agent.AgentType + tools []string + }{ + {"chat", "Chat Agent", "You are a helpful AI assistant. Provide clear and concise responses.", agent.AgentTypeChat, nil}, + {"react", "ReAct Agent", "You are a reasoning agent that can think and act. Break down complex problems step by step.", agent.AgentTypeReAct, []string{"calculator", "web_search"}}, + // Tool names must match the registry: "http" resolved to nothing, so + // the example agent advertised a tool it could never call. + {"tools", "Tool Agent", "You are a specialized agent that excels at using tools to accomplish tasks.", agent.AgentTypeTool, []string{"file_read", "file_write", "shell", "http_request"}}, } - fmt.Printf("βœ… Server stopped gracefully\n") - return nil -} + for _, example := range examples { + config := agent.DefaultAgentConfig() + config.ID = example.id + config.Name = example.name + config.Type = example.agentType + config.SystemPrompt = example.prompt + // DefaultAgentConfig leaves the provider and model empty, and the + // registry rejects a definition without a model. Registration therefore + // failed for all three agents while the command still reported + // "Created 3 example agents", leaving the server serving none of them. + config.Provider = "ollama" + config.Model = defaultExampleModel + if example.tools != nil { + config.Tools = example.tools + } -// createExampleAgents creates example agents if none are found -func createExampleAgents(autoServer *server.AutoServer) { - // Create a simple chat agent - chatConfig := agent.DefaultAgentConfig() - chatConfig.ID = "chat" - chatConfig.Name = "Chat Agent" - chatConfig.Type = agent.AgentTypeChat - chatConfig.SystemPrompt = "You are a helpful AI assistant. Provide clear and concise responses." - chatDefinition := agent.NewBaseAgentDefinition(chatConfig) - - autoServer.RegisterAgent("chat", chatDefinition) - - // Create a ReAct agent - reactConfig := agent.DefaultAgentConfig() - reactConfig.ID = "react" - reactConfig.Name = "ReAct Agent" - reactConfig.Type = agent.AgentTypeReAct - reactConfig.SystemPrompt = "You are a reasoning agent that can think and act. Break down complex problems step by step." - reactConfig.Tools = []string{"calculator", "web_search"} - reactDefinition := agent.NewBaseAgentDefinition(reactConfig) - - autoServer.RegisterAgent("react", reactDefinition) - - // Create a tool agent - toolConfig := agent.DefaultAgentConfig() - toolConfig.ID = "tools" - toolConfig.Name = "Tool Agent" - toolConfig.Type = agent.AgentTypeTool - toolConfig.SystemPrompt = "You are a specialized agent that excels at using tools to accomplish tasks." - toolConfig.Tools = []string{"file_read", "file_write", "shell", "http"} - toolDefinition := agent.NewBaseAgentDefinition(toolConfig) - - autoServer.RegisterAgent("tools", toolDefinition) - - fmt.Printf(" βœ… Created 3 example agents: chat, react, tools\n") + // Registration failures used to be printed and ignored, leaving the + // server serving fewer agents than it announced. + if err := autoServer.RegisterAgent(example.id, agent.NewBaseAgentDefinition(config)); err != nil { + return fmt.Errorf("failed to register the %s example agent: %w", example.id, err) + } + } + + _, _ = fmt.Fprintf(out, " βœ… Registered example agents: chat, react, tools\n") + return nil } // generateDockerfileForProject generates a Dockerfile for the project @@ -369,7 +495,7 @@ CMD ["./main", "auto-serve", "--host", "0.0.0.0", "--port", "8080"] ` dockerfilePath := filepath.Join(projectPath, "Dockerfile") - return os.WriteFile(dockerfilePath, []byte(dockerfileContent), 0644) + return writeFileChecked(dockerfilePath, dockerfileContent) } // generateDockerComposeForProject generates a docker-compose.yml @@ -415,7 +541,7 @@ volumes: `, config.Port, config.OllamaEndpoint) composePath := filepath.Join(projectPath, "docker-compose.yml") - return os.WriteFile(composePath, []byte(composeContent), 0644) + return writeFileChecked(composePath, composeContent) } // generateKubernetesManifests generates Kubernetes deployment manifests @@ -492,5 +618,5 @@ spec: `, config.OllamaEndpoint) manifestPath := filepath.Join(projectPath, "k8s-manifests.yaml") - return os.WriteFile(manifestPath, []byte(manifestContent), 0644) + return writeFileChecked(manifestPath, manifestContent) } diff --git a/cmd/golanggraph/auto_serve_command_test.go b/cmd/golanggraph/auto_serve_command_test.go new file mode 100644 index 0000000..aafe757 --- /dev/null +++ b/cmd/golanggraph/auto_serve_command_test.go @@ -0,0 +1,294 @@ +// Copyright (c) 2024 GoLangGraph Team +// +// Licensed under the MIT License. See LICENSE file in the project root for full license information. + +package main + +import ( + "bytes" + "fmt" + "os" + "path/filepath" + "strings" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/UnicoLab/GoLangGraph/pkg/agent" + "github.com/UnicoLab/GoLangGraph/pkg/server" + "github.com/UnicoLab/GoLangGraph/pkg/tools" +) + +// agentConfigYAMLFor returns a one-agent multi-agent config whose agent id is +// unique to the test. Agent definitions live in a process-wide registry, so +// fixtures must not collide between tests. +func agentConfigYAMLFor(t *testing.T) (id, yaml string) { + t.Helper() + + id = "agent-" + strings.ToLower(strings.NewReplacer("/", "-", "_", "-").Replace(t.Name())) + return id, fmt.Sprintf(`name: fixture +agents: + %s: + id: %s + name: Fixture Agent + type: chat + model: gpt-4 + provider: openai + systemprompt: "hi" + maxtokens: 1000 +`, id, id) +} + +// agentDirFor returns a directory holding one agent configuration unique to +// the test. Tests that do not want the built-in example agents registered as a +// side effect give the server a real agent to load instead. +func agentDirFor(t *testing.T) string { + t.Helper() + + _, yaml := agentConfigYAMLFor(t) + dir := t.TempDir() + writeTestFile(t, dir, "agents.yaml", yaml) + return dir +} + +func defaultAutoServeConfig(port int) *server.AutoServerConfig { + return &server.AutoServerConfig{ + Host: "127.0.0.1", + Port: port, + BasePath: "/api", + OllamaEndpoint: "http://localhost:11434", + ServerTimeout: 30 * time.Second, + MaxRequestSize: 1 << 20, + LLMProviders: map[string]interface{}{}, + } +} + +// Regression: --timeout and --max-request-size were declared on the command and +// then never read, so neither ever reached the server configuration. +func TestAutoServe_EveryDeclaredFlagReachesTheConfiguration(t *testing.T) { + t.Cleanup(func() { resetFlags(rootCmd) }) + + require.NoError(t, autoServeCmd.ParseFlags([]string{ + "--host", "127.0.0.1", + "--port", "9123", + "--base-path", "/v2", + "--timeout", "45s", + "--max-request-size", "2048", + "--ollama-endpoint", "http://ollama.internal:11434", + "--playground=false", + "--metrics=false", + "--env", "staging", + "--dev", + "--watch", + "--agent-dirs", "one,two", + "--openai-api-key", "sk-test", + })) + + config, opts, err := autoServeConfigFromFlags(autoServeCmd, []string{"./agents"}) + require.NoError(t, err) + + assert.Equal(t, "127.0.0.1", config.Host) + assert.Equal(t, 9123, config.Port) + assert.Equal(t, "/v2", config.BasePath) + assert.Equal(t, 45*time.Second, config.ServerTimeout, "--timeout must reach the server configuration") + assert.Equal(t, int64(2048), config.MaxRequestSize, "--max-request-size must reach the server configuration") + assert.Equal(t, "http://ollama.internal:11434", config.OllamaEndpoint) + assert.False(t, config.EnablePlayground) + assert.False(t, config.EnableMetricsAPI) + assert.Contains(t, config.LLMProviders, "openai") + + assert.Equal(t, "./agents", opts.SourcePath) + assert.Equal(t, "staging", opts.Env) + assert.True(t, opts.Dev) + assert.True(t, opts.Watch) + assert.Equal(t, []string{"one", "two"}, opts.AgentDirs) +} + +// Regression: a source path that did not exist was skipped in silence, and the +// command went on to serve three example agents while reporting that it had +// loaded agents from the operator's path. +func TestPrepareAutoServer_MissingSourcePathFails(t *testing.T) { + var out bytes.Buffer + _, err := prepareAutoServer(&out, defaultAutoServeConfig(8080), autoServeOptions{ + SourcePath: filepath.Join(t.TempDir(), "not-here"), + Env: "development", + }) + + require.Error(t, err) + assert.Contains(t, err.Error(), "agent source") +} + +// Regression: a file whose extension was not .yaml/.yml was ignored without a +// word, so "auto-serve agents.json" served nothing the operator had asked for. +func TestPrepareAutoServer_UnsupportedSourceExtensionFails(t *testing.T) { + dir := t.TempDir() + source := writeTestFile(t, dir, "agents.txt", "not a config") + + _, err := prepareAutoServer(&bytes.Buffer{}, defaultAutoServeConfig(8080), autoServeOptions{ + SourcePath: source, + Env: "development", + }) + + require.Error(t, err) + assert.Contains(t, err.Error(), "unsupported agent source") +} + +func TestPrepareAutoServer_LoadsAgentsFromAConfigFile(t *testing.T) { + id, yaml := agentConfigYAMLFor(t) + source := writeTestFile(t, t.TempDir(), "agents.yaml", yaml) + + var out bytes.Buffer + autoServer, err := prepareAutoServer(&out, defaultAutoServeConfig(8080), autoServeOptions{ + SourcePath: source, + Env: "development", + }) + + require.NoError(t, err) + require.NotNil(t, autoServer) + _, registered := agent.GetGlobalRegistry().GetDefinition(id) + assert.True(t, registered, "the agent in the config file must be registered") +} + +// server.AutoServer.LoadAgentsFromDirectory is a stub that loads nothing, so +// the CLI loads the configuration files in the directory itself. +func TestPrepareAutoServer_LoadsAgentsFromADirectory(t *testing.T) { + id, yaml := agentConfigYAMLFor(t) + dir := t.TempDir() + writeTestFile(t, dir, "agents.yaml", yaml) + writeTestFile(t, dir, "notes.md", "not a config") + + var out bytes.Buffer + _, err := prepareAutoServer(&out, defaultAutoServeConfig(8080), autoServeOptions{ + SourcePath: dir, + Env: "development", + }) + + require.NoError(t, err) + _, registered := agent.GetGlobalRegistry().GetDefinition(id) + assert.True(t, registered, "agents defined in the directory must be registered") + assert.Contains(t, out.String(), "1 agent config file(s) loaded") +} + +func TestPrepareAutoServer_RejectsInvalidPortsAndEnvironments(t *testing.T) { + _, err := prepareAutoServer(&bytes.Buffer{}, defaultAutoServeConfig(0), autoServeOptions{Env: "development"}) + require.Error(t, err) + + _, err = prepareAutoServer(&bytes.Buffer{}, defaultAutoServeConfig(8080), autoServeOptions{Env: "prod"}) + require.Error(t, err) + assert.Contains(t, err.Error(), "unknown environment") +} + +func TestPrepareAutoServer_ProductionDisablesThePlayground(t *testing.T) { + config := defaultAutoServeConfig(8080) + config.EnablePlayground = true + + _, err := prepareAutoServer(&bytes.Buffer{}, config, autoServeOptions{SourcePath: agentDirFor(t), Env: "production"}) + + require.NoError(t, err) + assert.False(t, config.EnablePlayground, "the playground must not be exposed in production") +} + +// Regression: --watch printed "πŸ‘€ File watching enabled (hot-reload)" next to a +// comment saying the watcher "would go here". Nothing ever watched anything. +func TestPrepareAutoServer_UnimplementedFlagsSaySo(t *testing.T) { + var out bytes.Buffer + _, err := prepareAutoServer(&out, defaultAutoServeConfig(8080), autoServeOptions{ + SourcePath: agentDirFor(t), + Env: "development", + Watch: true, + LogLevel: "debug", + }) + + require.NoError(t, err) + assert.Contains(t, out.String(), "--watch is not implemented") + assert.Contains(t, out.String(), "--log-level is not applied") + assert.NotContains(t, out.String(), "File watching enabled") +} + +// Regression: a plugin that failed to load was reported as a warning and the +// server started anyway, without the agents the operator asked for. +func TestPrepareAutoServer_MissingPluginFails(t *testing.T) { + dir := agentDirFor(t) + + _, err := prepareAutoServer(&bytes.Buffer{}, defaultAutoServeConfig(8080), autoServeOptions{ + SourcePath: dir, + Env: "development", + Plugins: []string{filepath.Join(dir, "absent.so")}, + }) + + require.Error(t, err) + assert.Contains(t, err.Error(), "failed to load plugin") +} + +func TestPrepareAutoServer_GeneratesDeploymentFiles(t *testing.T) { + dir := agentDirFor(t) + config := defaultAutoServeConfig(9090) + + var out bytes.Buffer + _, err := prepareAutoServer(&out, config, autoServeOptions{ + SourcePath: dir, + Env: "development", + GenerateDockerfile: true, + GenerateDockerCompose: true, + GenerateK8s: true, + }) + require.NoError(t, err) + + dockerfile, err := os.ReadFile(filepath.Join(dir, "Dockerfile")) + require.NoError(t, err) + assert.Contains(t, string(dockerfile), "FROM golang") + + compose, err := os.ReadFile(filepath.Join(dir, "docker-compose.yml")) + require.NoError(t, err) + assert.Contains(t, string(compose), "9090:8080", "the configured port must be used") + + manifests, err := os.ReadFile(filepath.Join(dir, "k8s-manifests.yaml")) + require.NoError(t, err) + assert.Contains(t, string(manifests), "kind: Deployment") +} + +func TestGenerateDeploymentFiles_ReportFailures(t *testing.T) { + missing := filepath.Join(t.TempDir(), "no-such-dir") + + require.Error(t, generateDockerfileForProject(missing)) + require.Error(t, generateDockerComposeForProject(missing, defaultAutoServeConfig(8080))) + require.Error(t, generateKubernetesManifests(missing, defaultAutoServeConfig(8080))) +} + +// The example agents are registered once per process; a failed registration +// used to be printed and ignored, leaving the server serving fewer agents than +// it announced. +func TestCreateExampleAgents_RegistersEveryAgentAndReportsCollisions(t *testing.T) { + autoServer := server.NewAutoServer(defaultAutoServeConfig(8080)) + + var out bytes.Buffer + require.NoError(t, createExampleAgents(&out, autoServer)) + + for _, id := range []string{"chat", "react", "tools"} { + _, registered := agent.GetGlobalRegistry().GetDefinition(id) + assert.True(t, registered, "example agent %q must be registered", id) + } + + err := createExampleAgents(&bytes.Buffer{}, autoServer) + require.Error(t, err, "a registration that failed must be reported, not printed and ignored") +} + +// Every tool an example agent asks for must exist, or the agent advertises a +// capability it can never use. +func TestExampleAgents_OnlyRequestRegisteredTools(t *testing.T) { + registry := agent.GetGlobalRegistry() + definition, ok := registry.GetDefinition("tools") + if !ok { + t.Skip("the example agents have not been registered in this run") + } + + known := map[string]bool{} + for _, name := range tools.NewToolRegistry().ListTools() { + known[name] = true + } + for _, tool := range definition.GetConfig().Tools { + assert.True(t, known[tool], "example agent requests unregistered tool %q", tool) + } +} diff --git a/cmd/golanggraph/cli_test.go b/cmd/golanggraph/cli_test.go new file mode 100644 index 0000000..9c5db0d --- /dev/null +++ b/cmd/golanggraph/cli_test.go @@ -0,0 +1,929 @@ +// Copyright (c) 2024 GoLangGraph Team +// +// Licensed under the MIT License. See LICENSE file in the project root for full license information. + +package main + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "go/parser" + "go/token" + "io" + "net" + "net/http" + "os" + "path/filepath" + "testing" + "time" + + "github.com/spf13/cobra" + "github.com/spf13/pflag" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// chdirTemp moves into a fresh temporary directory for the duration of a test. +// Commands such as init and docker build write relative to the working +// directory, and tests must not write outside their own temp dir. +func chdirTemp(t *testing.T) string { + t.Helper() + + dir := t.TempDir() + previous, err := os.Getwd() + require.NoError(t, err) + require.NoError(t, os.Chdir(dir)) + t.Cleanup(func() { + require.NoError(t, os.Chdir(previous)) + }) + return dir +} + +// writeTestFile writes a fixture file below dir. +func writeTestFile(t *testing.T, dir, name, content string) string { + t.Helper() + + path := filepath.Join(dir, name) + require.NoError(t, os.MkdirAll(filepath.Dir(path), 0o750)) + require.NoError(t, os.WriteFile(path, []byte(content), 0o600)) + return path +} + +// occupiedPort returns a port with a listener held open for the test, so a +// bind attempt against it is guaranteed to fail. +func occupiedPort(t *testing.T) int { + t.Helper() + + listener, err := net.Listen("tcp", "127.0.0.1:0") + require.NoError(t, err) + t.Cleanup(func() { _ = listener.Close() }) + return listener.Addr().(*net.TCPAddr).Port +} + +// freePort returns a port nothing is listening on. +func freePort(t *testing.T) int { + t.Helper() + + listener, err := net.Listen("tcp", "127.0.0.1:0") + require.NoError(t, err) + port := listener.Addr().(*net.TCPAddr).Port + require.NoError(t, listener.Close()) + return port +} + +const validAgentYAML = `name: "support-agent" +type: "chat" +model: "gpt-3.5-turbo" +provider: "openai" +system_prompt: "You are a helpful assistant." +temperature: 0.5 +max_tokens: 1000 +tools: + - calculator +` + +// --------------------------------------------------------------------------- +// validate +// --------------------------------------------------------------------------- + +// Regression: runValidate checked only that the file existed and then printed +// "Configuration validation completed successfully!" -- unparseable YAML was +// reported as a valid configuration. +func TestValidate_MalformedYAMLIsRejected(t *testing.T) { + dir := t.TempDir() + path := writeTestFile(t, dir, "broken.yaml", "this is not: [valid yaml at all\n - x\n") + + var out bytes.Buffer + err := runValidate(&out, []string{path}, false) + + require.Error(t, err, "malformed YAML must not be reported as valid") + assert.Contains(t, err.Error(), "parsing") + assert.NotContains(t, out.String(), "is valid") +} + +func TestValidate_MissingFileIsRejected(t *testing.T) { + var out bytes.Buffer + err := runValidate(&out, []string{filepath.Join(t.TempDir(), "absent.yaml")}, false) + + require.Error(t, err) + assert.NotContains(t, out.String(), "is valid") +} + +func TestValidate_EmptyFileIsRejected(t *testing.T) { + dir := t.TempDir() + path := writeTestFile(t, dir, "empty.yaml", " \n") + + err := runValidate(&bytes.Buffer{}, []string{path}, false) + require.Error(t, err) + assert.Contains(t, err.Error(), "empty") +} + +func TestValidate_UnsupportedExtensionIsRejected(t *testing.T) { + dir := t.TempDir() + path := writeTestFile(t, dir, "agent.txt", validAgentYAML) + + err := runValidate(&bytes.Buffer{}, []string{path}, false) + require.Error(t, err) + assert.Contains(t, err.Error(), "unsupported config extension") +} + +func TestValidate_AcceptsAValidConfiguration(t *testing.T) { + dir := t.TempDir() + path := writeTestFile(t, dir, "agent.yaml", validAgentYAML) + + var out bytes.Buffer + require.NoError(t, runValidate(&out, []string{path}, true)) + assert.Contains(t, out.String(), "is valid") +} + +func TestValidate_AcceptsJSON(t *testing.T) { + dir := t.TempDir() + path := writeTestFile(t, dir, "agent.json", `{ + "name": "json-agent", + "type": "chat", + "model": "gpt-4", + "provider": "openai", + "system_prompt": "hello", + "max_tokens": 500 + }`) + + require.NoError(t, runValidate(&bytes.Buffer{}, []string{path}, true)) +} + +func TestValidate_MissingRequiredFieldsAreReported(t *testing.T) { + dir := t.TempDir() + path := writeTestFile(t, dir, "agent.yaml", "name: incomplete\ntype: chat\n") + + var out bytes.Buffer + err := runValidate(&out, []string{path}, false) + + require.Error(t, err) + assert.Contains(t, out.String(), "model is required") +} + +// The framework's AgentConfig carries JSON tags only, so a YAML decode straight +// into it drops snake_case keys silently. A too-low max_tokens must therefore +// be seen and rejected; if the key were dropped the default of 1000 would +// quietly make this configuration "valid". +func TestValidate_SnakeCaseKeysAreActuallyRead(t *testing.T) { + dir := t.TempDir() + path := writeTestFile(t, dir, "agent.yaml", `name: low-tokens +type: chat +model: gpt-4 +provider: openai +system_prompt: "hi" +max_tokens: 50 +`) + + var out bytes.Buffer + err := runValidate(&out, []string{path}, false) + + require.Error(t, err, "max_tokens must be read from the file, not defaulted") + assert.Contains(t, out.String(), "MaxTokens too low") +} + +// The runtime silently falls back to a chat graph for an unknown agent type, +// so validation has to catch the typo. +func TestValidate_UnknownAgentTypeIsRejected(t *testing.T) { + dir := t.TempDir() + path := writeTestFile(t, dir, "agent.yaml", `name: typo +type: reactt +model: gpt-4 +provider: openai +max_tokens: 500 +`) + + var out bytes.Buffer + err := runValidate(&out, []string{path}, false) + + require.Error(t, err) + assert.Contains(t, out.String(), "unknown agent type") +} + +func TestValidate_StrictTurnsWarningsIntoFailures(t *testing.T) { + dir := t.TempDir() + path := writeTestFile(t, dir, "agent.yaml", `name: warns +type: chat +model: gpt-4 +provider: openai +system_prompt: "hi" +max_tokens: 500 +tools: + - no_such_tool +`) + + var out bytes.Buffer + require.NoError(t, runValidate(&out, []string{path}, false), "an unknown tool is only a warning") + assert.Contains(t, out.String(), "not registered") + + out.Reset() + err := runValidate(&out, []string{path}, true) + require.Error(t, err, "--strict must fail on warnings") + assert.Contains(t, err.Error(), "warning") +} + +func TestValidate_MultiAgentFileChecksRoutingTargets(t *testing.T) { + dir := t.TempDir() + path := writeTestFile(t, dir, "multi.yaml", `name: fleet +agents: + first: + name: First + type: chat + model: gpt-4 + provider: openai + system_prompt: "hi" + max_tokens: 500 +routing: + default_agent: first + rules: + - id: rule-1 + pattern: /first + agent_id: first + - id: rule-2 + pattern: /ghost + agent_id: ghost +`) + + var out bytes.Buffer + err := runValidate(&out, []string{path}, false) + + require.Error(t, err) + assert.Contains(t, out.String(), `rule targets agent "ghost"`) +} + +func TestValidate_ToolsMayBeObjects(t *testing.T) { + // The init template writes `tools: [{name: calculator, enabled: true}]`. + configs, err := loadAgentConfigs(writeTestFile(t, t.TempDir(), "agent.yaml", `name: obj-tools +type: chat +model: gpt-4 +provider: openai +max_tokens: 500 +tools: + - name: calculator + enabled: true + - name: web_search + enabled: false +`)) + require.NoError(t, err) + require.Len(t, configs, 1) + assert.Equal(t, []string{"calculator"}, configs[0].Tools, "disabled tools must not be requested") +} + +func TestValidate_RejectsWronglyTypedValues(t *testing.T) { + dir := t.TempDir() + path := writeTestFile(t, dir, "agent.yaml", "name: bad\ntype: chat\nmodel: gpt-4\nprovider: openai\ntemperature: hot\n") + + err := runValidate(&bytes.Buffer{}, []string{path}, false) + require.Error(t, err) + assert.Contains(t, err.Error(), "temperature") +} + +// --------------------------------------------------------------------------- +// init +// --------------------------------------------------------------------------- + +func TestInit_CreatesAProjectThatBuilds(t *testing.T) { + dir := chdirTemp(t) + + var out bytes.Buffer + require.NoError(t, runInit(&out, []string{"demo"}, "basic", false)) + + for _, name := range []string{"go.mod", "main.go", "README.md", ".gitignore", "docker-compose.yml", filepath.Join("configs", "agent-config.yaml")} { + assert.FileExists(t, filepath.Join(dir, "demo", name), "init must actually write %s", name) + } + + // The generated program has to be real Go, not a sketch. + source, err := os.ReadFile(filepath.Join(dir, "demo", "main.go")) + require.NoError(t, err) + _, err = parser.ParseFile(token.NewFileSet(), "main.go", source, parser.AllErrors) + require.NoError(t, err, "generated main.go must parse") + assert.Contains(t, string(source), "package main") + assert.Contains(t, string(source), "GoLangGraph/pkg/agent") + + gomod, err := os.ReadFile(filepath.Join(dir, "demo", "go.mod")) + require.NoError(t, err) + assert.Contains(t, string(gomod), "module demo") + + // And the configuration it ships with has to survive its own validator. + require.NoError(t, runValidate(&bytes.Buffer{}, []string{filepath.Join(dir, "demo", "configs", "agent-config.yaml")}, false)) +} + +// Regression: "golanggraph init ../escape" created and populated a directory +// outside the working directory. +func TestInit_RejectsNamesThatEscapeTheWorkingDirectory(t *testing.T) { + dir := chdirTemp(t) + parent := filepath.Dir(dir) + + for _, name := range []string{"../escape", "../../escape", "/tmp/escape-abs"} { + err := runInit(&bytes.Buffer{}, []string{name}, "basic", false) + require.Error(t, err, "init %q must be refused", name) + assert.NoDirExists(t, filepath.Join(parent, "escape")) + } + assert.NoDirExists(t, "/tmp/escape-abs") +} + +// Regression: an unknown template silently produced the basic project while +// reporting the template the operator had asked for. +func TestInit_RejectsUnknownTemplate(t *testing.T) { + dir := chdirTemp(t) + + err := runInit(&bytes.Buffer{}, []string{"demo"}, "nonsense", false) + require.Error(t, err) + assert.Contains(t, err.Error(), "unknown template") + assert.NoDirExists(t, filepath.Join(dir, "demo")) +} + +func TestInit_DoesNotOverwriteAnExistingProject(t *testing.T) { + dir := chdirTemp(t) + require.NoError(t, runInit(&bytes.Buffer{}, []string{"demo"}, "basic", false)) + + marker := writeTestFile(t, dir, filepath.Join("demo", "keep.txt"), "precious") + + err := runInit(&bytes.Buffer{}, []string{"demo"}, "basic", false) + require.Error(t, err) + assert.Contains(t, err.Error(), "--force") + + content, err := os.ReadFile(marker) + require.NoError(t, err) + assert.Equal(t, "precious", string(content)) + + require.NoError(t, runInit(&bytes.Buffer{}, []string{"demo"}, "basic", true), "--force must proceed") +} + +func TestInit_TemplatesWriteTheirOwnConfigs(t *testing.T) { + for _, tc := range []struct { + template string + expected []string + }{ + {"basic", []string{"configs/agent-config.yaml"}}, + {"advanced", []string{"configs/agent-config.yaml", "configs/advanced-config.yaml"}}, + {"rag", []string{"configs/agent-config.yaml", "configs/advanced-config.yaml", "configs/rag-config.yaml"}}, + } { + t.Run(tc.template, func(t *testing.T) { + dir := chdirTemp(t) + require.NoError(t, runInit(&bytes.Buffer{}, []string{"demo"}, tc.template, false)) + for _, rel := range tc.expected { + assert.FileExists(t, filepath.Join(dir, "demo", filepath.FromSlash(rel))) + } + }) + } +} + +func TestInit_ReportsWriteFailures(t *testing.T) { + dir := chdirTemp(t) + + // A file where the project directory should go makes every write fail. + writeTestFile(t, dir, "demo", "in the way") + + err := runInit(&bytes.Buffer{}, []string{"demo"}, "basic", true) + require.Error(t, err, "a failing scaffold must not report success") +} + +// --------------------------------------------------------------------------- +// debug visualize +// --------------------------------------------------------------------------- + +const graphYAML = `name: my-pipeline +start_node: ingest +end_nodes: [publish] +nodes: + - id: ingest + name: Ingest + - id: publish + name: Publish +edges: + - from: ingest + to: publish +` + +// Regression: visualize accepted a graph file argument and rendered a +// hard-coded three-node sample graph regardless of it. +func TestVisualize_RendersTheGraphFileGiven(t *testing.T) { + dir := t.TempDir() + path := writeTestFile(t, dir, "graph.yaml", graphYAML) + + var out bytes.Buffer + require.NoError(t, runVisualize(&out, []string{path}, "mermaid", "")) + + rendered := out.String() + assert.Contains(t, rendered, "ingest") + assert.Contains(t, rendered, "publish") + assert.NotContains(t, rendered, "process", "the built-in sample graph must not be rendered instead") +} + +// Regression: the json branch produced the literal string "JSON output not +// implemented yet", wrote it to the output file and reported success. +func TestVisualize_JSONFormatProducesJSON(t *testing.T) { + dir := t.TempDir() + path := writeTestFile(t, dir, "graph.yaml", graphYAML) + output := filepath.Join(dir, "topology.json") + + var out bytes.Buffer + require.NoError(t, runVisualize(&out, []string{path}, "json", output)) + + raw, err := os.ReadFile(output) + require.NoError(t, err) + assert.NotContains(t, string(raw), "not implemented") + + var topology struct { + Nodes []struct { + ID string `json:"id"` + IsStartNode bool `json:"is_start_node"` + } `json:"nodes"` + Edges []struct { + From string `json:"from"` + To string `json:"to"` + } `json:"edges"` + } + require.NoError(t, json.Unmarshal(raw, &topology)) + require.Len(t, topology.Nodes, 2) + require.Len(t, topology.Edges, 1) + assert.Equal(t, "ingest", topology.Edges[0].From) +} + +func TestVisualize_UnknownFormatFailsBeforeWriting(t *testing.T) { + dir := t.TempDir() + path := writeTestFile(t, dir, "graph.yaml", graphYAML) + output := filepath.Join(dir, "out.txt") + + err := runVisualize(&bytes.Buffer{}, []string{path}, "svg", output) + + require.Error(t, err) + assert.NoFileExists(t, output, "no file may be written for a format that cannot be produced") +} + +func TestVisualize_MalformedGraphFilesAreRejected(t *testing.T) { + dir := t.TempDir() + + for name, content := range map[string]string{ + "broken.yaml": "nodes: [oh dear\n", + "nonodes.yaml": "name: empty\n", + "badedge.yaml": "nodes:\n - id: a\nedges:\n - from: a\n to: ghost\n", + "badstart.yaml": `nodes: + - id: a +start_node: ghost +`, + } { + t.Run(name, func(t *testing.T) { + path := writeTestFile(t, dir, name, content) + require.Error(t, runVisualize(&bytes.Buffer{}, []string{path}, "mermaid", "")) + }) + } +} + +func TestVisualize_WithoutAFileSaysItIsASample(t *testing.T) { + var out bytes.Buffer + require.NoError(t, runVisualize(&out, nil, "mermaid", "")) + assert.Contains(t, out.String(), "built-in sample graph") +} + +// --------------------------------------------------------------------------- +// test +// --------------------------------------------------------------------------- + +// Regression: runTests ignored its argument entirely, built one hard-coded +// agent and printed "All tests completed successfully". +func TestTestCommand_ChecksTheConfigurationGiven(t *testing.T) { + dir := t.TempDir() + good := writeTestFile(t, dir, "good.yaml", validAgentYAML) + bad := writeTestFile(t, dir, "bad.yaml", "name: broken\ntype: chat\nprovider: openai\n") + + var out bytes.Buffer + require.NoError(t, runTests(&out, []string{good})) + assert.Contains(t, out.String(), "support-agent") + + require.Error(t, runTests(&bytes.Buffer{}, []string{bad}), "an invalid configuration must fail the test command") + require.Error(t, runTests(&bytes.Buffer{}, []string{filepath.Join(dir, "absent.yaml")})) +} + +func TestTestCommand_SelfTestPasses(t *testing.T) { + var out bytes.Buffer + require.NoError(t, runTests(&out, nil)) + assert.Contains(t, out.String(), "self-test") +} + +// --------------------------------------------------------------------------- +// migrate +// --------------------------------------------------------------------------- + +// Regression: the migrate flags were declared on the command but read back +// through viper, which they were never bound to, so every value arrived empty +// and "golanggraph migrate --db-type postgres" died with +// "Unsupported database type: ". +func TestMigrate_FlagsReachTheCommand(t *testing.T) { + out, err := executeRootCommand(t, "migrate", "--db-type", "bogus", "--db-host", "db.example.com", "--db-port", "6000") + + require.Error(t, err) + assert.Contains(t, err.Error(), `"bogus"`, "the --db-type value must reach the command") + assert.Contains(t, out, "db.example.com:6000", "the --db-host and --db-port values must reach the command") +} + +func TestMigrate_UnsupportedTypeIsRejected(t *testing.T) { + err := runMigrations(&bytes.Buffer{}, migrateOptions{Type: "mysql", Host: "localhost", Port: 3306}) + require.Error(t, err) + assert.Contains(t, err.Error(), "unsupported database type") +} + +func TestMigrate_MissingHostIsRejected(t *testing.T) { + require.Error(t, runMigrations(&bytes.Buffer{}, migrateOptions{Type: "postgres", Port: 5432})) + require.Error(t, runMigrations(&bytes.Buffer{}, migrateOptions{Type: "postgres", Host: "localhost", Port: 0})) +} + +// A database that cannot be reached must fail: the operator is about to deploy +// against the schema this command claims to have created. +func TestMigrate_UnreachableDatabaseFails(t *testing.T) { + port := freePort(t) + + t.Run("postgres", func(t *testing.T) { + err := runMigrations(&bytes.Buffer{}, migrateOptions{ + Type: "postgres", Host: "127.0.0.1", Port: port, Database: "golanggraph", Username: "postgres", + }) + require.Error(t, err) + assert.Contains(t, err.Error(), "postgres migration failed") + }) + + t.Run("redis", func(t *testing.T) { + err := runMigrations(&bytes.Buffer{}, migrateOptions{Type: "redis", Host: "127.0.0.1", Port: port}) + require.Error(t, err) + assert.Contains(t, err.Error(), "redis check failed") + }) +} + +// --------------------------------------------------------------------------- +// docker build / deploy +// --------------------------------------------------------------------------- + +// stubRunCommand replaces the external command runner for a test and records +// the command lines it was asked to run. +func stubRunCommand(t *testing.T, err error) *[][]string { + t.Helper() + + var recorded [][]string + original := runCommand + runCommand = func(ctx context.Context, out io.Writer, name string, args ...string) error { + recorded = append(recorded, append([]string{name}, args...)) + return err + } + t.Cleanup(func() { runCommand = original }) + return &recorded +} + +// Regression: docker build printed the command it would have run and reported +// "Docker build command prepared", so a pipeline calling it produced no image +// and no error. +func TestDockerBuild_ActuallyRunsDocker(t *testing.T) { + dir := chdirTemp(t) + config := writeTestFile(t, dir, "agent-config.yaml", validAgentYAML) + recorded := stubRunCommand(t, nil) + + var out bytes.Buffer + require.NoError(t, runDockerBuild(context.Background(), &out, []string{config}, dockerBuildOptions{ + Tag: "demo:1", Platform: "linux/amd64", + })) + + require.Len(t, *recorded, 1, "docker must be executed") + argv := (*recorded)[0] + assert.Equal(t, "docker", argv[0]) + assert.Contains(t, argv, "build") + assert.Contains(t, argv, "demo:1") + assert.Contains(t, argv, "linux/amd64") + assert.FileExists(t, filepath.Join(dir, "Dockerfile.agent")) +} + +func TestDockerBuild_DryRunDoesNotRunDocker(t *testing.T) { + dir := chdirTemp(t) + config := writeTestFile(t, dir, "agent-config.yaml", validAgentYAML) + recorded := stubRunCommand(t, nil) + + var out bytes.Buffer + require.NoError(t, runDockerBuild(context.Background(), &out, []string{config}, dockerBuildOptions{DryRun: true})) + + assert.Empty(t, *recorded) + assert.Contains(t, out.String(), "Dry run") +} + +func TestDockerBuild_ReportsDockerFailure(t *testing.T) { + dir := chdirTemp(t) + config := writeTestFile(t, dir, "agent-config.yaml", validAgentYAML) + stubRunCommand(t, fmt.Errorf("exit status 1")) + + err := runDockerBuild(context.Background(), &bytes.Buffer{}, []string{config}, dockerBuildOptions{}) + require.Error(t, err) + assert.Contains(t, err.Error(), "docker build failed") +} + +func TestDockerBuild_DoesNotOverwriteAnExistingDockerfile(t *testing.T) { + dir := chdirTemp(t) + config := writeTestFile(t, dir, "agent-config.yaml", validAgentYAML) + writeTestFile(t, dir, "Dockerfile.agent", "FROM scratch\n# hand written\n") + stubRunCommand(t, nil) + + require.NoError(t, runDockerBuild(context.Background(), &bytes.Buffer{}, []string{config}, dockerBuildOptions{})) + + content, err := os.ReadFile(filepath.Join(dir, "Dockerfile.agent")) + require.NoError(t, err) + assert.Contains(t, string(content), "hand written") +} + +func TestDockerBuild_RejectsAnInvalidConfiguration(t *testing.T) { + dir := chdirTemp(t) + broken := writeTestFile(t, dir, "broken.yaml", "nope: [\n") + recorded := stubRunCommand(t, nil) + + require.Error(t, runDockerBuild(context.Background(), &bytes.Buffer{}, []string{broken}, dockerBuildOptions{})) + assert.Empty(t, *recorded, "docker must not be invoked for a configuration that does not parse") +} + +func TestDockerBuild_MissingCustomDockerfileIsReported(t *testing.T) { + dir := chdirTemp(t) + config := writeTestFile(t, dir, "agent-config.yaml", validAgentYAML) + stubRunCommand(t, nil) + + err := runDockerBuild(context.Background(), &bytes.Buffer{}, []string{config}, + dockerBuildOptions{Dockerfile: filepath.Join(dir, "absent.Dockerfile")}) + require.Error(t, err) +} + +func TestDockerBuild_DistrolessUsesItsOwnDockerfile(t *testing.T) { + dir := chdirTemp(t) + config := writeTestFile(t, dir, "agent-config.yaml", validAgentYAML) + recorded := stubRunCommand(t, nil) + + require.NoError(t, runDockerBuild(context.Background(), &bytes.Buffer{}, []string{config}, dockerBuildOptions{Distroless: true})) + + assert.FileExists(t, filepath.Join(dir, "Dockerfile.distroless")) + assert.Contains(t, (*recorded)[0], "Dockerfile.distroless") +} + +// Regression: "deploy docker" printed "Docker deployment completed for config: +// X!" for any argument, including a path that did not exist, having done +// nothing at all. +func TestDeployDocker_DoesNotClaimAnUndoneDeployment(t *testing.T) { + dir := t.TempDir() + + t.Run("missing config", func(t *testing.T) { + var out bytes.Buffer + err := runDeployDocker(&out, []string{filepath.Join(dir, "absent.yaml")}) + require.Error(t, err) + assert.NotContains(t, out.String(), "completed") + }) + + t.Run("valid config", func(t *testing.T) { + config := writeTestFile(t, dir, "agent.yaml", validAgentYAML) + var out bytes.Buffer + err := runDeployDocker(&out, []string{config}) + require.Error(t, err, "a deployment that did not happen must not exit zero") + assert.ErrorIs(t, err, errNotImplemented) + assert.NotContains(t, out.String(), "completed") + }) +} + +// --------------------------------------------------------------------------- +// serve / dev +// --------------------------------------------------------------------------- + +// Regression: the server was started in a goroutine whose only error handling +// was log.Fatalf while the caller had already printed "Server started on +// host:port", so a failure to bind was announced as a success. +func TestServer_ReportsABindFailure(t *testing.T) { + port := occupiedPort(t) + + var out bytes.Buffer + err := runServer(context.Background(), &out, serverOptions{Host: "127.0.0.1", Port: port}) + + require.Error(t, err) + assert.Contains(t, err.Error(), "cannot bind") + assert.NotContains(t, out.String(), "listening") +} + +func TestServer_RejectsAnInvalidPort(t *testing.T) { + require.Error(t, runServer(context.Background(), &bytes.Buffer{}, serverOptions{Host: "127.0.0.1", Port: 0})) + require.Error(t, runServer(context.Background(), &bytes.Buffer{}, serverOptions{Host: "127.0.0.1", Port: 70000})) +} + +func TestServer_ServesAndShutsDownCleanly(t *testing.T) { + port := freePort(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + var out bytes.Buffer + done := make(chan error, 1) + go func() { + done <- runServer(ctx, &out, serverOptions{Host: "127.0.0.1", Port: port, StaticDir: t.TempDir()}) + }() + + url := fmt.Sprintf("http://127.0.0.1:%d/api/v1/health", port) + client := &http.Client{Timeout: time.Second} + + var reached bool + for i := 0; i < 100 && !reached; i++ { + resp, err := client.Get(url) + if err == nil { + _ = resp.Body.Close() + reached = true + break + } + time.Sleep(20 * time.Millisecond) + } + require.True(t, reached, "the server must answer on the port it reported") + + cancel() + select { + case err := <-done: + require.NoError(t, err) + case <-time.After(35 * time.Second): + t.Fatal("the server did not shut down when the context was canceled") + } +} + +// Regression: the dev command declared its own --host/--port flags but +// runDevServer read host and port from viper, where only the serve command's +// flags were bound, so "golanggraph dev --port 3000" started on 8080. +func TestDev_UsesItsOwnPortFlag(t *testing.T) { + port := occupiedPort(t) + + _, err := executeRootCommand(t, "dev", "--host", "127.0.0.1", "--port", fmt.Sprint(port)) + + require.Error(t, err) + assert.Contains(t, err.Error(), fmt.Sprintf("127.0.0.1:%d", port), + "the dev command must bind the port its own flag names") +} + +func TestDev_LoadsTheAgentConfigFlag(t *testing.T) { + dir := t.TempDir() + broken := writeTestFile(t, dir, "broken.yaml", "agents: [\n") + + err := runServer(context.Background(), &bytes.Buffer{}, serverOptions{ + Host: "127.0.0.1", Port: freePort(t), Dev: true, AgentConfig: broken, + }) + + require.Error(t, err, "--agent-config must be read, not ignored") + assert.Contains(t, err.Error(), "agent config") +} + +// Regression: --log-level was declared on the dev command and never read. +func TestDev_ValidatesAndAppliesTheLogLevel(t *testing.T) { + err := runServer(context.Background(), &bytes.Buffer{}, serverOptions{ + Host: "127.0.0.1", Port: freePort(t), Dev: true, LogLevel: "shouting", + }) + + require.Error(t, err, "an unusable log level must be reported, not ignored") + assert.Contains(t, err.Error(), "invalid log level") +} + +func TestDev_SaysHotReloadIsNotImplemented(t *testing.T) { + // Point at an occupied port so the run stops right after the notice. + var out bytes.Buffer + _ = runServer(context.Background(), &out, serverOptions{ + Host: "127.0.0.1", Port: occupiedPort(t), Dev: true, HotReload: true, + }) + + assert.Contains(t, out.String(), "hot-reload is not implemented") +} + +// --------------------------------------------------------------------------- +// command wiring +// --------------------------------------------------------------------------- + +// executeRootCommand runs the real cobra command tree with the given arguments +// and returns its output. Flag values persist on the shared command objects, so +// the flags are restored afterwards. +func executeRootCommand(t *testing.T, args ...string) (string, error) { + t.Helper() + + var out bytes.Buffer + rootCmd.SetOut(&out) + rootCmd.SetErr(&out) + rootCmd.SetArgs(args) + t.Cleanup(func() { + rootCmd.SetOut(nil) + rootCmd.SetErr(nil) + rootCmd.SetArgs(nil) + resetFlags(rootCmd) + }) + + err := rootCmd.Execute() + return out.String(), err +} + +// resetFlags restores every flag in the tree to its default value. +func resetFlags(cmd *cobra.Command) { + cmd.Flags().VisitAll(func(f *pflag.Flag) { + if f.Changed { + _ = f.Value.Set(f.DefValue) + f.Changed = false + } + }) + for _, sub := range cmd.Commands() { + resetFlags(sub) + } +} + +// Every leaf command must do something: a command with neither a Run nor +// subcommands only prints its own help. +func TestRootCommand_EveryLeafCommandIsRunnable(t *testing.T) { + var walk func(cmd *cobra.Command, path string) + walk = func(cmd *cobra.Command, path string) { + if len(cmd.Commands()) == 0 { + assert.True(t, cmd.Run != nil || cmd.RunE != nil, "%s has nothing to run", path) + return + } + for _, sub := range cmd.Commands() { + walk(sub, path+" "+sub.Name()) + } + } + walk(rootCmd, rootCmd.Name()) +} + +func TestRootCommand_KnownCommandsArePresent(t *testing.T) { + names := map[string]bool{} + for _, cmd := range rootCmd.Commands() { + names[cmd.Name()] = true + } + for _, expected := range []string{ + "auto-serve", "debug", "deploy", "dev", "docker", "health", + "init", "migrate", "multi-agent", "serve", "test", "validate", + } { + assert.True(t, names[expected], "command %q is missing", expected) + } +} + +func TestSafeProjectDir(t *testing.T) { + for _, name := range []string{"", " ", "..", "../x", "/abs"} { + _, err := safeProjectDir(name) + assert.Error(t, err, "%q must be refused", name) + } + for input, expected := range map[string]string{ + "demo": "demo", + "./demo": "demo", + "nested/app": filepath.Join("nested", "app"), + } { + got, err := safeProjectDir(input) + require.NoError(t, err) + assert.Equal(t, expected, got) + } +} + +func TestCheckAddressAvailable(t *testing.T) { + require.NoError(t, checkAddressAvailable("127.0.0.1", freePort(t))) + + err := checkAddressAvailable("127.0.0.1", occupiedPort(t)) + require.Error(t, err) + assert.Contains(t, err.Error(), "cannot bind") +} + +func TestWriteFileChecked_ReportsFailure(t *testing.T) { + dir := t.TempDir() + require.NoError(t, writeFileChecked(filepath.Join(dir, "ok.txt"), "content")) + + err := writeFileChecked(filepath.Join(dir, "ok.txt", "nested.txt"), "content") + require.Error(t, err) + assert.Contains(t, err.Error(), "failed to write") +} + +func TestNormalizeKey(t *testing.T) { + for _, spelling := range []string{"system_prompt", "systemPrompt", "system-prompt", "SystemPrompt"} { + assert.Equal(t, "systemprompt", normalizeKey(spelling)) + } +} + +func TestLoadAgentConfigs_MultiAgentFilesAreOrdered(t *testing.T) { + path := writeTestFile(t, t.TempDir(), "multi.yaml", `name: fleet +agents: + zeta: + name: Zeta + type: chat + model: m + provider: ollama + alpha: + name: Alpha + type: chat + model: m + provider: ollama +`) + + configs, err := loadAgentConfigs(path) + require.NoError(t, err) + require.Len(t, configs, 2) + assert.Equal(t, "alpha", configs[0].Key, "agents must be reported in a stable order") + assert.Equal(t, "zeta", configs[1].Key) +} + +func TestLoadAgentConfigs_RejectsMalformedAgentsSection(t *testing.T) { + dir := t.TempDir() + for name, content := range map[string]string{ + "list.yaml": "name: fleet\nagents:\n - not-a-map\n", + "empty.yaml": "name: fleet\nagents: {}\n", + "scalar.yaml": `name: fleet +agents: + one: "just a string" +`, + } { + t.Run(name, func(t *testing.T) { + _, err := loadAgentConfigs(writeTestFile(t, dir, name, content)) + require.Error(t, err) + }) + } +} diff --git a/cmd/golanggraph/config_yaml_regression_test.go b/cmd/golanggraph/config_yaml_regression_test.go new file mode 100644 index 0000000..de33e09 --- /dev/null +++ b/cmd/golanggraph/config_yaml_regression_test.go @@ -0,0 +1,54 @@ +// Copyright (c) 2024 GoLangGraph Team +// +// Licensed under the MIT License. See LICENSE file in the project root for full license information. + +package main + +import ( + "context" + "testing" + + "github.com/stretchr/testify/require" + "gopkg.in/yaml.v3" + + "github.com/UnicoLab/GoLangGraph/pkg/agent" + "github.com/UnicoLab/GoLangGraph/pkg/core" + "github.com/UnicoLab/GoLangGraph/pkg/llm" +) + +// Config structs are written to disk as YAML by `golanggraph multi-agent init`. +// A func-typed field carries no YAML representation, and gopkg.in/yaml.v3 +// panics rather than skipping it: adding EarlyExit to AgentConfig with only a +// `json:"-"` tag turned project scaffolding into a crash. Every func-typed +// field in a serialisable config needs `yaml:"-"` as well. +func TestConfigStructs_MarshalToYAMLWithoutPanicking(t *testing.T) { + exit := func(content string, calls []llm.ToolCall) bool { return true } + + for name, value := range map[string]interface{}{ + "agent.AgentConfig": &agent.AgentConfig{ + ID: "a1", + Name: "scaffolded", + Type: agent.AgentTypeChat, + EarlyExit: exit, + }, + "llm.CompletionRequest": &llm.CompletionRequest{ + Model: "m", + EarlyExit: exit, + }, + "core.RetryPolicy": &core.RetryPolicy{ + MaxAttempts: 2, + RetryIf: func(error) bool { return false }, + }, + "core.Node": &core.Node{ + ID: "n", + Function: func(ctx context.Context, s *core.BaseState) (*core.BaseState, error) { return s, nil }, + }, + } { + t.Run(name, func(t *testing.T) { + out, err := yaml.Marshal(value) + require.NoError(t, err, "%s must marshal to YAML", name) + require.NotContains(t, string(out), "earlyexit", + "a func field must be omitted, not emitted") + }) + } +} diff --git a/cmd/golanggraph/health.go b/cmd/golanggraph/health.go new file mode 100644 index 0000000..b7a997c --- /dev/null +++ b/cmd/golanggraph/health.go @@ -0,0 +1,280 @@ +// Copyright (c) 2024 GoLangGraph Team +// +// Licensed under the MIT License. See LICENSE file in the project root for full license information. +// +// Package: GoLangGraph - A powerful Go framework for building AI agent workflows + +package main + +import ( + "context" + "fmt" + "net" + "net/http" + "net/url" + "os" + "strings" + "syscall" + "time" +) + +// healthOptions configures a health check run. +type healthOptions struct { + // ServerURL, when set, probes a running server's HTTP health endpoint + // instead of inspecting local dependencies. This is what a container + // health check should use. + ServerURL string + // Strict turns warnings into failures. + Strict bool + // Timeout bounds each individual probe. + Timeout time.Duration + // MinFreeDiskBytes is the free space below which the check fails. + MinFreeDiskBytes uint64 +} + +// checkResult is the outcome of a single probe. +type checkResult struct { + Name string + OK bool + Warning bool + Detail string +} + +func (r checkResult) symbol() string { + switch { + case r.OK: + return "βœ“" + case r.Warning: + return "⚠" + default: + return "βœ—" + } +} + +// runHealthCheck performs a real health check and exits with a status an +// orchestrator can act on. +// +// The previous implementation printed "βœ“ Reachable" for PostgreSQL, Redis and +// Ollama without ever connecting to them, and reported disk and memory as fine +// unconditionally β€” so it reported a healthy system no matter what was actually +// running. It also exited non-zero when the optional OPENAI_API_KEY was unset, +// which combined with the container HEALTHCHECK marked every deployment without +// an OpenAI key permanently unhealthy. +func runHealthCheck(opts healthOptions) int { + if opts.Timeout <= 0 { + opts.Timeout = 3 * time.Second + } + if opts.MinFreeDiskBytes == 0 { + opts.MinFreeDiskBytes = 64 << 20 // 64 MiB + } + + fmt.Printf("Running GoLangGraph health check...\n") + + var results []checkResult + + if opts.ServerURL != "" { + results = append(results, probeServer(opts.ServerURL, opts.Timeout)) + } else { + results = append(results, probeDependencies(opts)...) + results = append(results, probeResources(opts)...) + } + + failed, warned := 0, 0 + for _, r := range results { + fmt.Printf(" %s %s: %s\n", r.symbol(), r.Name, r.Detail) + switch { + case !r.OK && !r.Warning: + failed++ + case r.Warning: + warned++ + } + } + + fmt.Printf("\n") + switch { + case failed > 0: + fmt.Printf("βœ— System is unhealthy: %d failed, %d warnings\n", failed, warned) + return 1 + case warned > 0 && opts.Strict: + fmt.Printf("βœ— System has %d warnings and --strict is set\n", warned) + return 1 + case warned > 0: + fmt.Printf("βœ… System is healthy (%d warnings)\n", warned) + return 0 + default: + fmt.Printf("βœ… System is healthy\n") + return 0 + } +} + +// probeServer checks a running server's health endpoint. +func probeServer(rawURL string, timeout time.Duration) checkResult { + name := "server " + rawURL + + parsed, err := url.Parse(rawURL) + if err != nil || parsed.Host == "" { + return checkResult{Name: name, Detail: fmt.Sprintf("invalid server URL: %v", err)} + } + endpoint := strings.TrimRight(rawURL, "/") + "/api/v1/health" + + ctx, cancel := context.WithTimeout(context.Background(), timeout) + defer cancel() + + req, err := http.NewRequestWithContext(ctx, http.MethodGet, endpoint, nil) + if err != nil { + return checkResult{Name: name, Detail: err.Error()} + } + + resp, err := http.DefaultClient.Do(req) + if err != nil { + return checkResult{Name: name, Detail: fmt.Sprintf("unreachable: %v", err)} + } + defer func() { _ = resp.Body.Close() }() + + if resp.StatusCode != http.StatusOK { + return checkResult{Name: name, Detail: fmt.Sprintf("returned status %d", resp.StatusCode)} + } + return checkResult{Name: name, OK: true, Detail: "responding"} +} + +// probeDependencies checks the services that are actually configured. +// +// A dependency is only probed when its environment variable is set: defaulting +// to localhost would report a failure in every deployment that does not use +// that service. +func probeDependencies(opts healthOptions) []checkResult { + var results []checkResult + + if host := os.Getenv("POSTGRES_HOST"); host != "" { + port := envOr("POSTGRES_PORT", "5432") + results = append(results, probeTCP("PostgreSQL", net.JoinHostPort(host, port), opts.Timeout, false)) + } + + if host := os.Getenv("REDIS_HOST"); host != "" { + port := envOr("REDIS_PORT", "6379") + results = append(results, probeTCP("Redis", net.JoinHostPort(host, port), opts.Timeout, false)) + } + + if raw := os.Getenv("OLLAMA_URL"); raw != "" { + if parsed, err := url.Parse(raw); err == nil && parsed.Host != "" { + host := parsed.Hostname() + port := parsed.Port() + if port == "" { + port = map[bool]string{true: "443", false: "80"}[parsed.Scheme == "https"] + } + results = append(results, probeTCP("Ollama", net.JoinHostPort(host, port), opts.Timeout, false)) + } else { + results = append(results, checkResult{Name: "Ollama", Detail: "invalid OLLAMA_URL"}) + } + } + + // Credentials are informational: a deployment may legitimately use only one + // provider, so a missing key is never a failure on its own. + for _, cred := range []struct{ name, env string }{ + {"OpenAI credentials", "OPENAI_API_KEY"}, + {"Gemini credentials", "GEMINI_API_KEY"}, + } { + if os.Getenv(cred.env) != "" { + results = append(results, checkResult{Name: cred.name, OK: true, Detail: "configured"}) + } else { + results = append(results, checkResult{Name: cred.name, Warning: true, Detail: cred.env + " is not set"}) + } + } + + if len(results) == 0 { + results = append(results, checkResult{ + Name: "dependencies", OK: true, + Detail: "none configured; set POSTGRES_HOST, REDIS_HOST or OLLAMA_URL to probe them", + }) + } + return results +} + +// probeTCP opens a TCP connection to verify a service is actually listening. +func probeTCP(name, address string, timeout time.Duration, optional bool) checkResult { + conn, err := net.DialTimeout("tcp", address, timeout) + if err != nil { + return checkResult{ + Name: name, + Warning: optional, + Detail: fmt.Sprintf("%s unreachable: %v", address, err), + } + } + _ = conn.Close() + return checkResult{Name: name, OK: true, Detail: address + " reachable"} +} + +// probeResources measures real disk and memory availability. +func probeResources(opts healthOptions) []checkResult { + var results []checkResult + + wd, err := os.Getwd() + if err != nil { + wd = "." + } + + var stat syscall.Statfs_t + if err := syscall.Statfs(wd, &stat); err != nil { + results = append(results, checkResult{Name: "Disk space", Warning: true, Detail: "unavailable: " + err.Error()}) + } else { + free := stat.Bavail * uint64(stat.Bsize) + detail := fmt.Sprintf("%s free at %s", humanBytes(free), wd) + if free < opts.MinFreeDiskBytes { + results = append(results, checkResult{Name: "Disk space", Detail: detail + " (below the minimum)"}) + } else { + results = append(results, checkResult{Name: "Disk space", OK: true, Detail: detail}) + } + } + + results = append(results, probeMemory()) + return results +} + +// probeMemory reports available memory from the kernel where it is exposed. +func probeMemory() checkResult { + data, err := os.ReadFile("/proc/meminfo") + if err != nil { + return checkResult{Name: "Memory", Warning: true, Detail: "unavailable on this platform"} + } + + for _, line := range strings.Split(string(data), "\n") { + if !strings.HasPrefix(line, "MemAvailable:") { + continue + } + fields := strings.Fields(line) + if len(fields) < 2 { + break + } + var kb uint64 + if _, err := fmt.Sscanf(fields[1], "%d", &kb); err != nil { + break + } + available := kb * 1024 + detail := humanBytes(available) + " available" + if available < 64<<20 { + return checkResult{Name: "Memory", Detail: detail + " (low)"} + } + return checkResult{Name: "Memory", OK: true, Detail: detail} + } + return checkResult{Name: "Memory", Warning: true, Detail: "MemAvailable not reported"} +} + +func envOr(key, fallback string) string { + if v := os.Getenv(key); v != "" { + return v + } + return fallback +} + +func humanBytes(n uint64) string { + const unit = 1024 + if n < unit { + return fmt.Sprintf("%d B", n) + } + div, exp := uint64(unit), 0 + for v := n / unit; v >= unit && exp < 4; v /= unit { + div *= unit + exp++ + } + return fmt.Sprintf("%.1f %ciB", float64(n)/float64(div), "KMGTP"[exp]) +} diff --git a/cmd/golanggraph/health_test.go b/cmd/golanggraph/health_test.go new file mode 100644 index 0000000..98288a5 --- /dev/null +++ b/cmd/golanggraph/health_test.go @@ -0,0 +1,160 @@ +// Copyright (c) 2024 GoLangGraph Team +// +// Licensed under the MIT License. See LICENSE file in the project root for full license information. + +package main + +import ( + "net" + "net/http" + "net/http/httptest" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// A container with no dependencies configured is healthy. The previous +// implementation exited non-zero whenever OPENAI_API_KEY was unset, which made +// every deployment without an OpenAI key permanently unhealthy. +func TestHealth_NoDependenciesConfiguredIsHealthy(t *testing.T) { + t.Setenv("POSTGRES_HOST", "") + t.Setenv("REDIS_HOST", "") + t.Setenv("OLLAMA_URL", "") + t.Setenv("OPENAI_API_KEY", "") + t.Setenv("GEMINI_API_KEY", "") + + code := runHealthCheck(healthOptions{Timeout: time.Second}) + assert.Equal(t, 0, code, "a missing optional credential must not fail the check") +} + +// With --strict, warnings do fail. +func TestHealth_StrictTreatsWarningsAsFailures(t *testing.T) { + t.Setenv("POSTGRES_HOST", "") + t.Setenv("REDIS_HOST", "") + t.Setenv("OLLAMA_URL", "") + t.Setenv("OPENAI_API_KEY", "") + t.Setenv("GEMINI_API_KEY", "") + + code := runHealthCheck(healthOptions{Timeout: time.Second, Strict: true}) + assert.Equal(t, 1, code) +} + +// A configured dependency that is not listening must fail. The previous +// implementation printed "Reachable" without connecting to anything. +func TestHealth_UnreachableDependencyFails(t *testing.T) { + // Bind and release a port so nothing is listening on it. + listener, err := net.Listen("tcp", "127.0.0.1:0") + require.NoError(t, err) + port := listener.Addr().(*net.TCPAddr).Port + require.NoError(t, listener.Close()) + + t.Setenv("POSTGRES_HOST", "127.0.0.1") + t.Setenv("POSTGRES_PORT", itoa(port)) + t.Setenv("REDIS_HOST", "") + t.Setenv("OLLAMA_URL", "") + + code := runHealthCheck(healthOptions{Timeout: time.Second}) + assert.Equal(t, 1, code, "an unreachable configured dependency must be reported") +} + +// A dependency that is listening passes. +func TestHealth_ReachableDependencyPasses(t *testing.T) { + listener, err := net.Listen("tcp", "127.0.0.1:0") + require.NoError(t, err) + defer func() { _ = listener.Close() }() + go func() { + for { + conn, err := listener.Accept() + if err != nil { + return + } + _ = conn.Close() + } + }() + + port := listener.Addr().(*net.TCPAddr).Port + t.Setenv("POSTGRES_HOST", "127.0.0.1") + t.Setenv("POSTGRES_PORT", itoa(port)) + t.Setenv("REDIS_HOST", "") + t.Setenv("OLLAMA_URL", "") + + code := runHealthCheck(healthOptions{Timeout: 2 * time.Second}) + assert.Equal(t, 0, code) +} + +// The server probe reports what the endpoint says, which is what the container +// health check depends on. +func TestHealth_ServerProbe(t *testing.T) { + t.Run("healthy", func(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/api/v1/health" { + w.WriteHeader(http.StatusNotFound) + return + } + w.WriteHeader(http.StatusOK) + })) + defer srv.Close() + + assert.Equal(t, 0, runHealthCheck(healthOptions{ServerURL: srv.URL, Timeout: 2 * time.Second})) + }) + + t.Run("erroring", func(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusServiceUnavailable) + })) + defer srv.Close() + + assert.Equal(t, 1, runHealthCheck(healthOptions{ServerURL: srv.URL, Timeout: 2 * time.Second})) + }) + + t.Run("unreachable", func(t *testing.T) { + listener, err := net.Listen("tcp", "127.0.0.1:0") + require.NoError(t, err) + addr := listener.Addr().String() + require.NoError(t, listener.Close()) + + assert.Equal(t, 1, runHealthCheck(healthOptions{ + ServerURL: "http://" + addr, Timeout: time.Second, + })) + }) + + t.Run("invalid URL", func(t *testing.T) { + assert.Equal(t, 1, runHealthCheck(healthOptions{ServerURL: "://nonsense", Timeout: time.Second})) + }) +} + +// Resource probes must report real measurements, not fixed strings. +func TestHealth_ResourceProbesAreMeasured(t *testing.T) { + results := probeResources(healthOptions{MinFreeDiskBytes: 1}) + require.NotEmpty(t, results) + + for _, r := range results { + assert.NotEmpty(t, r.Detail) + assert.NotEqual(t, "βœ“ Sufficient", r.Detail) + assert.NotEqual(t, "βœ“ Available", r.Detail) + } + + // An impossible free-space requirement must fail rather than pass anyway. + strictResults := probeResources(healthOptions{MinFreeDiskBytes: 1 << 62}) + var sawFailure bool + for _, r := range strictResults { + if r.Name == "Disk space" && !r.OK && !r.Warning { + sawFailure = true + } + } + assert.True(t, sawFailure, "disk space must be measured, not assumed") +} + +func itoa(n int) string { + if n == 0 { + return "0" + } + var digits []byte + for n > 0 { + digits = append([]byte{byte('0' + n%10)}, digits...) + n /= 10 + } + return string(digits) +} diff --git a/cmd/golanggraph/main.go b/cmd/golanggraph/main.go index bbfcdae..c7ee90d 100644 --- a/cmd/golanggraph/main.go +++ b/cmd/golanggraph/main.go @@ -8,15 +8,24 @@ package main import ( "context" + "encoding/json" + "errors" "fmt" - "log" + "io" + "net" "os" + "os/exec" "os/signal" + "path/filepath" + "sort" + "strings" "syscall" "time" + "github.com/sirupsen/logrus" "github.com/spf13/cobra" "github.com/spf13/viper" + yaml "gopkg.in/yaml.v3" "github.com/UnicoLab/GoLangGraph/pkg/agent" "github.com/UnicoLab/GoLangGraph/pkg/core" @@ -27,6 +36,12 @@ import ( "github.com/UnicoLab/GoLangGraph/pkg/tools" ) +// errNotImplemented marks a command that cannot do what its name claims. Such a +// command must fail loudly: several commands here used to print a success +// message ("Docker deployment completed", "Deploying to Docker...") and exit 0 +// without doing any work at all, which an operator would act on. +var errNotImplemented = errors.New("not implemented") + var ( cfgFile string verbose bool @@ -41,11 +56,11 @@ for building stateful, multi-agent conversational AI applications. This CLI provides tools for: - Building and packaging agents into Docker containers -- Running development servers with hot-reload +- Running a development server - Managing database migrations - Visualizing graph execution -- Testing and debugging agents -- Deploying agents to production environments`, +- Validating and testing agent configurations +- Generating deployment artifacts`, } // serveCmd represents the serve command @@ -58,8 +73,15 @@ The server provides: - WebSocket endpoints for real-time streaming - Visual debugging interface - Health monitoring endpoints`, - Run: func(cmd *cobra.Command, args []string) { - runServer() + RunE: func(cmd *cobra.Command, args []string) error { + ctx, stop := signal.NotifyContext(cmd.Context(), syscall.SIGINT, syscall.SIGTERM) + defer stop() + return runServer(ctx, cmd.OutOrStdout(), serverOptions{ + Host: viper.GetString("host"), + Port: viper.GetInt("port"), + StaticDir: viper.GetString("static-dir"), + CORS: viper.GetBool("enable-cors"), + }) }, } @@ -67,9 +89,41 @@ The server provides: var migrateCmd = &cobra.Command{ Use: "migrate", Short: "Run database migrations", - Long: `Run database migrations to set up the required schema for state persistence.`, - Run: func(cmd *cobra.Command, args []string) { - runMigrations() + Long: `Run database migrations to set up the required schema for state persistence. + +For postgres this creates the threads, checkpoints, sessions and document tables +if they do not exist. Redis has no schema; for redis this only verifies that the +server is reachable.`, + RunE: func(cmd *cobra.Command, args []string) error { + opts := migrateOptions{} + var err error + // Defect: these flags were declared on this command but the values were + // read back through viper, which they were never bound to. Every value + // came back empty, so "golanggraph migrate --db-host db.example.com" + // ignored the host entirely and "--db-type postgres" reached the switch + // as "" and died with `Unsupported database type: `. Read the flags. + if opts.Type, err = cmd.Flags().GetString("db-type"); err != nil { + return err + } + if opts.Host, err = cmd.Flags().GetString("db-host"); err != nil { + return err + } + if opts.Port, err = cmd.Flags().GetInt("db-port"); err != nil { + return err + } + if opts.Database, err = cmd.Flags().GetString("db-name"); err != nil { + return err + } + if opts.Username, err = cmd.Flags().GetString("db-user"); err != nil { + return err + } + if opts.Password, err = cmd.Flags().GetString("db-password"); err != nil { + return err + } + if opts.SSLMode, err = cmd.Flags().GetString("db-sslmode"); err != nil { + return err + } + return runMigrations(cmd.OutOrStdout(), opts) }, } @@ -84,22 +138,50 @@ var debugCmd = &cobra.Command{ var visualizeCmd = &cobra.Command{ Use: "visualize [graph-file]", Short: "Visualize a graph structure", - Long: `Generate visual representations of graph structures in various formats (Mermaid, DOT, JSON).`, - Args: cobra.MaximumNArgs(1), - Run: func(cmd *cobra.Command, args []string) { - format, _ := cmd.Flags().GetString("format") - output, _ := cmd.Flags().GetString("output") - runVisualize(args, format, output) + Long: `Generate visual representations of graph structures in various formats (mermaid, dot, json). + +The graph file is a JSON or YAML document describing the graph: + + name: my-graph + start_node: start + end_nodes: [finish] + nodes: + - id: start + name: Start + - id: finish + name: Finish + edges: + - from: start + to: finish + +Without a graph file a built-in sample graph is rendered, and the output says so.`, + Args: cobra.MaximumNArgs(1), + RunE: func(cmd *cobra.Command, args []string) error { + format, err := cmd.Flags().GetString("format") + if err != nil { + return err + } + output, err := cmd.Flags().GetString("output") + if err != nil { + return err + } + return runVisualize(cmd.OutOrStdout(), args, format, output) }, } // testCmd represents the test command var testCmd = &cobra.Command{ - Use: "test", + Use: "test [config-file]", Short: "Test agent configurations and graph execution", - Long: `Test agent configurations and validate graph execution flows.`, - Run: func(cmd *cobra.Command, args []string) { - runTests() + Long: `Test an agent configuration by building the agent it describes and validating +the execution graph that results. + +With a configuration file the agent in that file is built and checked. Without +one, a built-in self-test exercises agent construction and graph validation so +that a broken installation is detected. No LLM calls are made either way.`, + Args: cobra.MaximumNArgs(1), + RunE: func(cmd *cobra.Command, args []string) error { + return runTests(cmd.OutOrStdout(), args) }, } @@ -107,18 +189,47 @@ var testCmd = &cobra.Command{ var healthCmd = &cobra.Command{ Use: "health", Short: "Check system health and component status", - Long: `Check the health status of GoLangGraph components including databases, LLM providers, and system resources.`, + Long: `Check the health status of GoLangGraph components. + +With --server, probes a running server's HTTP health endpoint; this is what a +container health check should use. Otherwise it probes the dependencies that +are actually configured (POSTGRES_HOST, REDIS_HOST, OLLAMA_URL) and the local +disk and memory. + +Missing optional provider credentials are reported as warnings and do not fail +the check unless --strict is given.`, Run: func(cmd *cobra.Command, args []string) { - runHealthCheck() + serverURL, _ := cmd.Flags().GetString("server") + strict, _ := cmd.Flags().GetBool("strict") + timeout, _ := cmd.Flags().GetDuration("timeout") + + if serverURL == "" { + serverURL = os.Getenv("GOLANGGRAPH_SERVER_URL") + } + + os.Exit(runHealthCheck(healthOptions{ + ServerURL: serverURL, + Strict: strict, + Timeout: timeout, + })) }, } -// buildCmd represents the build command +// buildCmd represents the build command. +// +// It has no subcommands of its own: invoking it used to print the help text and +// exit 0, which reads as "the build succeeded". Point at the command that does +// the work and fail. var buildCmd = &cobra.Command{ Use: "build", Short: "Build and package agents for deployment", Long: `Build and package agents into deployable artifacts including Docker containers. -Supports both regular and distroless container builds for production deployment.`, + +Container images are built by "golanggraph docker build", which supports both +regular and distroless variants.`, + RunE: func(cmd *cobra.Command, args []string) error { + return fmt.Errorf("nothing to build here: use 'golanggraph docker build [agent-config]'") + }, } // dockerCmd represents the docker command @@ -135,27 +246,71 @@ var dockerBuildCmd = &cobra.Command{ Long: `Build a Docker container for deploying an agent to production. Supports both regular and distroless variants for different deployment needs.`, Args: cobra.MaximumNArgs(1), - Run: func(cmd *cobra.Command, args []string) { - distroless, _ := cmd.Flags().GetBool("distroless") - tag, _ := cmd.Flags().GetString("tag") - dockerfile, _ := cmd.Flags().GetString("dockerfile") - platform, _ := cmd.Flags().GetString("platform") - runDockerBuild(args, distroless, tag, dockerfile, platform) + RunE: func(cmd *cobra.Command, args []string) error { + opts := dockerBuildOptions{} + var err error + if opts.Distroless, err = cmd.Flags().GetBool("distroless"); err != nil { + return err + } + if opts.Tag, err = cmd.Flags().GetString("tag"); err != nil { + return err + } + if opts.Dockerfile, err = cmd.Flags().GetString("dockerfile"); err != nil { + return err + } + if opts.Platform, err = cmd.Flags().GetString("platform"); err != nil { + return err + } + if opts.DryRun, err = cmd.Flags().GetBool("dry-run"); err != nil { + return err + } + if opts.ContextDir, err = cmd.Flags().GetString("context"); err != nil { + return err + } + return runDockerBuild(cmd.Context(), cmd.OutOrStdout(), args, opts) }, } // devCmd represents the dev command var devCmd = &cobra.Command{ Use: "dev", - Short: "Start development server with hot-reload", - Long: `Start a development server with hot-reload capabilities for testing and debugging agents. -Includes: -- Hot-reload on code changes -- Interactive debugging interface -- Real-time logging and metrics -- Agent testing playground`, - Run: func(cmd *cobra.Command, args []string) { - runDevServer() + Short: "Start a development server", + Long: `Start a development server for testing and debugging agents. + +Includes an interactive debugging interface and an agent playground. Hot-reload +is not implemented: --hot-reload only reports that, so restart the server to +pick up changes.`, + RunE: func(cmd *cobra.Command, args []string) error { + ctx, stop := signal.NotifyContext(cmd.Context(), syscall.SIGINT, syscall.SIGTERM) + defer stop() + + // Defect: dev declares its own --host/--port/--hot-reload/--log-level + // flags, but runDevServer read host and port from viper, where only the + // *serve* command's flags are bound. "golanggraph dev --port 3000" + // therefore started on 8080. Read this command's own flags. + opts := serverOptions{Dev: true, CORS: true, StaticDir: "./static"} + var err error + if opts.Host, err = cmd.Flags().GetString("host"); err != nil { + return err + } + if opts.Port, err = cmd.Flags().GetInt("port"); err != nil { + return err + } + if opts.AgentConfig, err = cmd.Flags().GetString("agent-config"); err != nil { + return err + } + if opts.HotReload, err = cmd.Flags().GetBool("hot-reload"); err != nil { + return err + } + if opts.LogLevel, err = cmd.Flags().GetString("log-level"); err != nil { + return err + } + debug, err := cmd.Flags().GetBool("debug") + if err != nil { + return err + } + opts.Dev = debug + return runServer(ctx, cmd.OutOrStdout(), opts) }, } @@ -163,11 +318,22 @@ Includes: var validateCmd = &cobra.Command{ Use: "validate [config-file]", Short: "Validate agent configuration", - Long: `Validate agent configuration files and graph definitions for correctness.`, - Args: cobra.MaximumNArgs(1), - Run: func(cmd *cobra.Command, args []string) { - strict, _ := cmd.Flags().GetBool("strict") - runValidate(args, strict) + Long: `Validate agent configuration files and graph definitions for correctness. + +Both single-agent files and multi-agent files (a top-level "agents:" map) are +understood, in YAML or JSON. The file is parsed, required fields are checked, +value ranges are checked, tool names are resolved against the built-in tool +registry and the resulting agent graph is built and validated. + +With --strict, warnings (unknown keys, missing system prompt, unknown provider) +are treated as errors.`, + Args: cobra.MaximumNArgs(1), + RunE: func(cmd *cobra.Command, args []string) error { + strict, err := cmd.Flags().GetBool("strict") + if err != nil { + return err + } + return runValidate(cmd.OutOrStdout(), args, strict) }, } @@ -182,10 +348,15 @@ var deployCmd = &cobra.Command{ var deployDockerCmd = &cobra.Command{ Use: "docker [agent-config]", Short: "Deploy agent using Docker", - Long: `Deploy an agent using Docker containers with production-ready configuration.`, - Args: cobra.MaximumNArgs(1), - Run: func(cmd *cobra.Command, args []string) { - runDeployDocker(args) + Long: `Deploy an agent using Docker containers with production-ready configuration. + +This command validates the agent configuration and then reports that pushing and +running the container is not implemented; it does not pretend to have deployed +anything. Build an image with "golanggraph docker build" and run it with docker +or docker compose.`, + Args: cobra.MaximumNArgs(1), + RunE: func(cmd *cobra.Command, args []string) error { + return runDeployDocker(cmd.OutOrStdout(), args) }, } @@ -193,17 +364,33 @@ var deployDockerCmd = &cobra.Command{ var initCmd = &cobra.Command{ Use: "init [project-name]", Short: "Initialize a new GoLangGraph project", - Long: `Initialize a new GoLangGraph project with example configurations and templates.`, - Args: cobra.MaximumNArgs(1), - Run: func(cmd *cobra.Command, args []string) { - template, _ := cmd.Flags().GetString("template") - runInit(args, template) + Long: `Initialize a new GoLangGraph project with example configurations and templates. + +The project name is used as a directory name below the current directory; names +that escape it (absolute paths, "..") are rejected. An existing non-empty +directory is not overwritten unless --force is given.`, + Args: cobra.MaximumNArgs(1), + RunE: func(cmd *cobra.Command, args []string) error { + template, err := cmd.Flags().GetString("template") + if err != nil { + return err + } + force, err := cmd.Flags().GetBool("force") + if err != nil { + return err + } + return runInit(cmd.OutOrStdout(), args, template, force) }, } func init() { cobra.OnInitialize(initConfig) + // A command that fails at runtime should report the failure, not bury it + // under a page of usage text, and main() prints the error itself. + rootCmd.SilenceUsage = true + rootCmd.SilenceErrors = true + // Global flags rootCmd.PersistentFlags().StringVar(&cfgFile, "config", "", "config file (default is $HOME/.golanggraph.yaml)") rootCmd.PersistentFlags().BoolVarP(&verbose, "verbose", "v", false, "verbose output") @@ -218,8 +405,8 @@ func init() { devCmd.Flags().StringP("host", "H", "localhost", "Host to bind to") devCmd.Flags().IntP("port", "p", 8080, "Port to bind to") devCmd.Flags().String("agent-config", "", "Agent configuration file") - devCmd.Flags().Bool("hot-reload", true, "Enable hot-reload") - devCmd.Flags().Bool("debug", true, "Enable debug mode") + devCmd.Flags().Bool("hot-reload", true, "Enable hot-reload (not implemented)") + devCmd.Flags().Bool("debug", true, "Enable the server's development mode (debug interface and playground)") devCmd.Flags().String("log-level", "info", "Log level (debug, info, warn, error)") // Docker build command flags @@ -227,12 +414,15 @@ func init() { dockerBuildCmd.Flags().StringP("tag", "t", "", "Docker image tag") dockerBuildCmd.Flags().String("dockerfile", "", "Custom Dockerfile path") dockerBuildCmd.Flags().String("platform", "", "Target platform (e.g., linux/amd64,linux/arm64)") + dockerBuildCmd.Flags().Bool("dry-run", false, "Print the docker command without running it") + dockerBuildCmd.Flags().String("context", ".", "Docker build context directory") // Validate command flags validateCmd.Flags().BoolP("strict", "s", false, "Enable strict validation") // Init command flags initCmd.Flags().StringP("template", "t", "basic", "Project template (basic, advanced, rag)") + initCmd.Flags().Bool("force", false, "Overwrite an existing project directory") // Migrate command flags migrateCmd.Flags().String("db-type", "postgres", "Database type (postgres, redis)") @@ -241,6 +431,7 @@ func init() { migrateCmd.Flags().String("db-name", "golanggraph", "Database name") migrateCmd.Flags().String("db-user", "postgres", "Database user") migrateCmd.Flags().String("db-password", "", "Database password") + migrateCmd.Flags().String("db-sslmode", "disable", "PostgreSQL sslmode (disable, require, verify-full)") // Visualize command flags visualizeCmd.Flags().StringP("format", "f", "mermaid", "Output format (mermaid, dot, json)") @@ -257,6 +448,9 @@ func init() { rootCmd.AddCommand(migrateCmd) rootCmd.AddCommand(debugCmd) rootCmd.AddCommand(testCmd) + healthCmd.Flags().String("server", "", "probe a running server's health endpoint instead of local dependencies") + healthCmd.Flags().Bool("strict", false, "treat warnings as failures") + healthCmd.Flags().Duration("timeout", 3*time.Second, "per-probe timeout") rootCmd.AddCommand(healthCmd) // Add nested commands @@ -265,10 +459,10 @@ func init() { debugCmd.AddCommand(visualizeCmd) // Bind flags to viper - viper.BindPFlag("host", serveCmd.Flags().Lookup("host")) - viper.BindPFlag("port", serveCmd.Flags().Lookup("port")) - viper.BindPFlag("static-dir", serveCmd.Flags().Lookup("static-dir")) - viper.BindPFlag("enable-cors", serveCmd.Flags().Lookup("enable-cors")) + _ = viper.BindPFlag("host", serveCmd.Flags().Lookup("host")) + _ = viper.BindPFlag("port", serveCmd.Flags().Lookup("port")) + _ = viper.BindPFlag("static-dir", serveCmd.Flags().Lookup("static-dir")) + _ = viper.BindPFlag("enable-cors", serveCmd.Flags().Lookup("enable-cors")) } // initConfig reads in config file and ENV variables. @@ -292,181 +486,381 @@ func initConfig() { // If a config file is found, read it in. if err := viper.ReadInConfig(); err == nil && verbose { - fmt.Fprintln(os.Stderr, "Using config file:", viper.ConfigFileUsed()) + _, _ = fmt.Fprintln(os.Stderr, "Using config file:", viper.ConfigFileUsed()) } } -func runServer() { - fmt.Println("Starting GoLangGraph server...") +// serverOptions describes a serve or dev run. +type serverOptions struct { + Host string + Port int + StaticDir string + CORS bool + // Dev enables the development mode of the server. + Dev bool + // AgentConfig is an optional agent configuration file whose agents are + // created on the server before it starts serving. + AgentConfig string + // HotReload is the dev command's --hot-reload flag. File watching is not + // implemented; the flag is reported honestly rather than acted on. + HotReload bool + // LogLevel is applied to the server's logger. + LogLevel string +} + +// runServer starts the HTTP server and blocks until ctx is canceled. +// +// Two defects are fixed here. The server used to be started in a goroutine +// whose only error handling was log.Fatalf, while the caller had already +// printed "Server started on host:port" -- so a failure to bind was announced +// as a success. And the dev command's own flags were ignored (see devCmd). +func runServer(ctx context.Context, out io.Writer, opts serverOptions) error { + if opts.Port <= 0 || opts.Port > 65535 { + return fmt.Errorf("invalid port %d", opts.Port) + } + if opts.Host == "" { + opts.Host = "0.0.0.0" + } + if opts.StaticDir == "" { + opts.StaticDir = "./static" + } + + // --log-level was declared on the dev command and never read. + if opts.LogLevel != "" { + if _, err := logrus.ParseLevel(opts.LogLevel); err != nil { + return fmt.Errorf("invalid log level %q (want debug, info, warn or error)", opts.LogLevel) + } + } - // Create server configuration config := &server.ServerConfig{ - Host: viper.GetString("host"), - Port: viper.GetInt("port"), + Host: opts.Host, + Port: opts.Port, ReadTimeout: 30 * time.Second, WriteTimeout: 30 * time.Second, MaxHeaderBytes: 1 << 20, - EnableCORS: viper.GetBool("enable-cors"), - StaticDir: viper.GetString("static-dir"), + EnableCORS: opts.CORS, + StaticDir: opts.StaticDir, + DevMode: opts.Dev, + LogLevel: opts.LogLevel, + } + + if opts.Dev { + _, _ = fmt.Fprintln(out, "Starting GoLangGraph development server...") + } else { + _, _ = fmt.Fprintln(out, "Starting GoLangGraph server...") } - // Create server srv := server.NewServer(config) - // Initialize components - if err := initializeComponents(srv); err != nil { - log.Fatalf("Failed to initialize components: %v", err) + agentManager, err := initializeComponents(out, srv) + if err != nil { + return fmt.Errorf("failed to initialize components: %w", err) } - // Start server in a goroutine - go func() { - if err := srv.Start(); err != nil { - log.Fatalf("Server failed to start: %v", err) + // --agent-config used to be declared and never read. Load it for real. + if opts.AgentConfig != "" { + configs, err := loadAgentConfigs(opts.AgentConfig) + if err != nil { + return fmt.Errorf("agent config %s: %w", opts.AgentConfig, err) + } + for _, cfg := range configs { + if _, err := agentManager.CreateAgent(cfg.toAgentConfig()); err != nil { + return fmt.Errorf("failed to create agent %s: %w", cfg.Name, err) + } + _, _ = fmt.Fprintf(out, "Loaded agent %s (%s)\n", cfg.Name, cfg.Type) } - }() + } - fmt.Printf("Server started on %s:%d\n", config.Host, config.Port) - fmt.Printf("Health check: http://%s:%d/api/v1/health\n", config.Host, config.Port) + if opts.HotReload { + // The previous implementation printed "Hot-reload enabled - watching + // for changes..." next to a comment saying the watcher "would go here". + _, _ = fmt.Fprintln(out, "Note: hot-reload is not implemented; restart the server to pick up changes") + } - // Wait for interrupt signal to gracefully shutdown - quit := make(chan os.Signal, 1) - signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM) - <-quit + // Fail before announcing success if the address cannot be bound. + if err := checkAddressAvailable(opts.Host, opts.Port); err != nil { + return err + } + + errCh := make(chan error, 1) + go func() { errCh <- srv.Start() }() - fmt.Println("Shutting down server...") + _, _ = fmt.Fprintf(out, "Server listening on %s:%d\n", config.Host, config.Port) + _, _ = fmt.Fprintf(out, "Health check: http://%s:%d/api/v1/health\n", config.Host, config.Port) + if opts.Dev { + _, _ = fmt.Fprintf(out, "Debug interface: http://%s:%d/debug\n", config.Host, config.Port) + _, _ = fmt.Fprintf(out, "Agent playground: http://%s:%d/playground\n", config.Host, config.Port) + } - // Create a deadline for shutdown - ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + select { + case err := <-errCh: + if err != nil { + return fmt.Errorf("server failed to start: %w", err) + } + return nil + case <-ctx.Done(): + } + + _, _ = fmt.Fprintln(out, "Shutting down server...") + + shutdownCtx, cancel := context.WithTimeout(context.Background(), 30*time.Second) defer cancel() - if err := srv.Stop(ctx); err != nil { - log.Fatalf("Server forced to shutdown: %v", err) + if err := srv.Stop(shutdownCtx); err != nil { + return fmt.Errorf("server forced to shutdown: %w", err) } - fmt.Println("Server exited") + _, _ = fmt.Fprintln(out, "Server exited") + return nil } -func initializeComponents(srv *server.Server) error { - // Initialize LLM providers +// checkAddressAvailable reports whether the server can bind host:port. Without +// this the CLI announced a listening server whose bind had already failed. +func checkAddressAvailable(host string, port int) error { + listener, err := net.Listen("tcp", net.JoinHostPort(host, fmt.Sprintf("%d", port))) + if err != nil { + return fmt.Errorf("cannot bind %s:%d: %w", host, port, err) + } + return listener.Close() +} + +// initializeComponents wires the LLM providers, tools and managers onto the +// server and returns the agent manager it installed. +// +// Provider construction errors used to be discarded with `if err == nil`, so a +// misconfigured provider silently disappeared, and the built-in tools were +// re-registered on a registry that already contains them, discarding the +// "already registered" errors that came back. +func initializeComponents(out io.Writer, srv *server.Server) (*server.AgentManager, error) { llmManager := llm.NewProviderManager() - // Add OpenAI provider if API key is available if apiKey := os.Getenv("OPENAI_API_KEY"); apiKey != "" { - openaiConfig := &llm.ProviderConfig{ + openaiProvider, err := llm.NewOpenAIProvider(&llm.ProviderConfig{ APIKey: apiKey, Endpoint: "https://api.openai.com/v1", + }) + if err != nil { + return nil, fmt.Errorf("openai provider: %w", err) } - openaiProvider, err := llm.NewOpenAIProvider(openaiConfig) - if err == nil { - llmManager.RegisterProvider("openai", openaiProvider) + if err := llmManager.RegisterProvider("openai", openaiProvider); err != nil { + return nil, fmt.Errorf("register openai provider: %w", err) } } - // Add Ollama provider if available - if ollamaURL := os.Getenv("OLLAMA_URL"); ollamaURL != "" { - ollamaConfig := &llm.ProviderConfig{ - Endpoint: ollamaURL, - } - ollamaProvider, err := llm.NewOllamaProvider(ollamaConfig) - if err == nil { - llmManager.RegisterProvider("ollama", ollamaProvider) - } - } else { - // Default Ollama URL - ollamaConfig := &llm.ProviderConfig{ - Endpoint: "http://localhost:11434", - } - ollamaProvider, err := llm.NewOllamaProvider(ollamaConfig) - if err == nil { - llmManager.RegisterProvider("ollama", ollamaProvider) - } + ollamaURL := os.Getenv("OLLAMA_URL") + if ollamaURL == "" { + ollamaURL = "http://localhost:11434" + } + ollamaProvider, err := llm.NewOllamaProvider(&llm.ProviderConfig{Endpoint: ollamaURL}) + if err != nil { + return nil, fmt.Errorf("ollama provider: %w", err) + } + if err := llmManager.RegisterProvider("ollama", ollamaProvider); err != nil { + return nil, fmt.Errorf("register ollama provider: %w", err) } - // Initialize tool registry + // NewToolRegistry already registers the built-in tools. toolRegistry := tools.NewToolRegistry() - // Register default tools - toolRegistry.RegisterTool(tools.NewWebSearchTool()) - toolRegistry.RegisterTool(tools.NewCalculatorTool()) - toolRegistry.RegisterTool(tools.NewFileReadTool()) - toolRegistry.RegisterTool(tools.NewFileWriteTool()) - toolRegistry.RegisterTool(tools.NewShellTool()) - toolRegistry.RegisterTool(tools.NewHTTPTool()) - toolRegistry.RegisterTool(tools.NewTimeTool()) - - // Initialize session manager (using memory for now) sessionManager := persistence.NewSessionManager(nil) - - // Initialize agent manager agentManager := server.NewAgentManager(llmManager, toolRegistry) - // Set components on server srv.SetLLMManager(llmManager) srv.SetToolRegistry(toolRegistry) srv.SetAgentManager(agentManager) srv.SetSessionManager(sessionManager) - return nil + _, _ = fmt.Fprintf(out, "Providers: %s | Tools: %d\n", + strings.Join(llmManager.ListProviders(), ", "), len(toolRegistry.ListTools())) + + return agentManager, nil } -func runMigrations() { - fmt.Println("Running database migrations...") +// migrateOptions describes a migrate run. +type migrateOptions struct { + Type string + Host string + Port int + Database string + Username string + Password string + SSLMode string +} + +func runMigrations(out io.Writer, opts migrateOptions) error { + if opts.Host == "" { + return errors.New("database host is required (--db-host)") + } + if opts.Port <= 0 { + return fmt.Errorf("invalid database port %d", opts.Port) + } - dbType := viper.GetString("db-type") + _, _ = fmt.Fprintf(out, "Running database migrations against %s %s:%d...\n", opts.Type, opts.Host, opts.Port) - switch dbType { - case "postgres": + switch opts.Type { + case "postgres", "postgresql", "pgvector": + sslMode := opts.SSLMode + if sslMode == "" { + sslMode = "disable" + } config := &persistence.DatabaseConfig{ - Type: "postgres", - Host: viper.GetString("db-host"), - Port: viper.GetInt("db-port"), - Database: viper.GetString("db-name"), - Username: viper.GetString("db-user"), - Password: viper.GetString("db-password"), - SSLMode: "disable", + Type: persistence.DatabaseType(opts.Type), + Host: opts.Host, + Port: opts.Port, + Database: opts.Database, + Username: opts.Username, + Password: opts.Password, + SSLMode: sslMode, + // Without a connect timeout a wrong host hangs for minutes. + ConnectionParams: map[string]string{"connect_timeout": "10"}, } + // NewPostgresCheckpointer runs the schema migration (CREATE TABLE IF + // NOT EXISTS ...) as part of construction. checkpointer, err := persistence.NewPostgresCheckpointer(config) if err != nil { - log.Fatalf("Failed to create PostgreSQL checkpointer: %v", err) + return fmt.Errorf("postgres migration failed: %w", err) } - defer checkpointer.Close() + defer func() { + if cerr := checkpointer.Close(); cerr != nil { + _, _ = fmt.Fprintf(out, "warning: failed to close database connection: %v\n", cerr) + } + }() - fmt.Println("PostgreSQL migrations completed successfully") + _, _ = fmt.Fprintln(out, "PostgreSQL schema is up to date (threads, checkpoints, sessions, documents)") + return nil case "redis": config := &persistence.DatabaseConfig{ - Type: "redis", - Host: viper.GetString("db-host"), - Port: viper.GetInt("db-port"), - Password: viper.GetString("db-password"), + Type: persistence.DatabaseTypeRedis, + Host: opts.Host, + Port: opts.Port, + Password: opts.Password, } checkpointer, err := persistence.NewRedisCheckpointer(config) if err != nil { - log.Fatalf("Failed to create Redis checkpointer: %v", err) + return fmt.Errorf("redis check failed: %w", err) } - defer checkpointer.Close() + defer func() { + if cerr := checkpointer.Close(); cerr != nil { + _, _ = fmt.Fprintf(out, "warning: failed to close redis connection: %v\n", cerr) + } + }() - fmt.Println("Redis setup completed successfully") + // Redis is schemaless: there is nothing to migrate. Saying "Redis setup + // completed successfully" implied work that never happened. + _, _ = fmt.Fprintln(out, "Redis has no schema to migrate; the server is reachable") + return nil default: - log.Fatalf("Unsupported database type: %s", dbType) + return fmt.Errorf("unsupported database type: %q (want postgres, pgvector or redis)", opts.Type) } } -func runVisualize(args []string, format, output string) { - fmt.Printf("Visualizing graph in %s format...\n", format) +// graphFile is the on-disk description of a graph to visualize. +type graphFile struct { + Name string `json:"name" yaml:"name"` + StartNode string `json:"start_node" yaml:"start_node"` + EndNodes []string `json:"end_nodes" yaml:"end_nodes"` + Nodes []graphFileNode `json:"nodes" yaml:"nodes"` + Edges []graphFileEdge `json:"edges" yaml:"edges"` +} - // Create a sample graph for demonstration - // In a real implementation, this would load from a file or configuration - sampleGraph := createSampleGraph() +type graphFileNode struct { + ID string `json:"id" yaml:"id"` + Name string `json:"name" yaml:"name"` +} - // Create visualizer - visualizer := debug.NewGraphVisualizer(nil, nil) +type graphFileEdge struct { + From string `json:"from" yaml:"from"` + To string `json:"to" yaml:"to"` +} + +// loadGraphFromFile builds a graph from a JSON or YAML description. +// +// The visualize command used to accept a graph file argument and then render a +// hard-coded three-node sample graph regardless of it, so an operator inspected +// a diagram that had nothing to do with their graph. +func loadGraphFromFile(path string) (*core.Graph, error) { + data, err := os.ReadFile(path) // #nosec G304 -- operator-supplied path is the point of the command + if err != nil { + return nil, err + } + + var spec graphFile + switch ext := strings.ToLower(filepath.Ext(path)); ext { + case ".json": + if err := json.Unmarshal(data, &spec); err != nil { + return nil, fmt.Errorf("parsing %s: %w", path, err) + } + case ".yaml", ".yml": + if err := yaml.Unmarshal(data, &spec); err != nil { + return nil, fmt.Errorf("parsing %s: %w", path, err) + } + default: + return nil, fmt.Errorf("unsupported graph file extension %q (want .json, .yaml or .yml)", ext) + } + + if len(spec.Nodes) == 0 { + return nil, fmt.Errorf("%s defines no nodes", path) + } - // Get topology - topology := visualizer.GetGraphTopology(sampleGraph) + name := spec.Name + if name == "" { + name = strings.TrimSuffix(filepath.Base(path), filepath.Ext(path)) + } + + graph := core.NewGraph(name) + for _, node := range spec.Nodes { + nodeName := node.Name + if nodeName == "" { + nodeName = node.ID + } + // Visualization only needs the topology, so every node gets an + // identity function. + graph.AddNode(node.ID, nodeName, func(ctx context.Context, state *core.BaseState) (*core.BaseState, error) { + return state, nil + }) + } + for _, edge := range spec.Edges { + graph.AddEdge(edge.From, edge.To, nil) + } + + start := spec.StartNode + if start == "" { + start = spec.Nodes[0].ID + } + if err := graph.SetStartNode(start); err != nil { + return nil, err + } + for _, end := range spec.EndNodes { + if err := graph.AddEndNode(end); err != nil { + return nil, err + } + } + + if err := graph.Validate(); err != nil { + return nil, fmt.Errorf("graph in %s is invalid: %w", path, err) + } + return graph, nil +} + +func runVisualize(out io.Writer, args []string, format, output string) error { + graph := createSampleGraph() + source := "built-in sample graph" + + if len(args) > 0 { + loaded, err := loadGraphFromFile(args[0]) + if err != nil { + return err + } + graph = loaded + source = args[0] + } + + visualizer := debug.NewGraphVisualizer(nil, nil) + topology := visualizer.GetGraphTopology(graph) var result string switch format { @@ -475,25 +869,37 @@ func runVisualize(args []string, format, output string) { case "dot": result = visualizer.GenerateDotDiagram(topology) case "json": - // JSON output would need to be implemented - result = "JSON output not implemented yet" + // This branch used to emit the literal string "JSON output not + // implemented yet" -- and then write it to the output file and report + // "Visualization saved to ". + encoded, err := json.MarshalIndent(topology, "", " ") + if err != nil { + return fmt.Errorf("failed to encode topology: %w", err) + } + result = string(encoded) default: - log.Fatalf("Unsupported format: %s", format) + return fmt.Errorf("unsupported format %q (want mermaid, dot or json)", format) + } + + _, _ = fmt.Fprintf(out, "Visualizing %s in %s format...\n", source, format) + if len(args) == 0 { + _, _ = fmt.Fprintln(out, "(no graph file given; pass one to visualize your own graph)") } - // Output result if output != "" { if err := os.WriteFile(output, []byte(result), 0600); err != nil { - log.Fatalf("Failed to write output file: %v", err) + return fmt.Errorf("failed to write output file: %w", err) } - fmt.Printf("Visualization saved to %s\n", output) - } else { - fmt.Println(result) + _, _ = fmt.Fprintf(out, "Visualization saved to %s\n", output) + return nil } + + _, _ = fmt.Fprintln(out, result) + return nil } +// createSampleGraph builds the graph rendered when no graph file is given. func createSampleGraph() *core.Graph { - // This is a placeholder - in a real implementation, you'd load from configuration graph := core.NewGraph("sample-graph") // Add some sample nodes @@ -513,249 +919,854 @@ func createSampleGraph() *core.Graph { graph.AddEdge("start", "process", nil) graph.AddEdge("process", "end", nil) - // Set start and end nodes - graph.SetStartNode("start") - graph.AddEndNode("end") + // Set start and end nodes. The nodes were added immediately above, so these + // cannot fail; Validate() would surface it if they ever did. + _ = graph.SetStartNode("start") + _ = graph.AddEndNode("end") return graph } -func runTests() { - fmt.Println("Running tests...") +// runTests builds the agents described by a configuration file (or a built-in +// one) and validates the graphs they produce. +// +// The command used to ignore any argument, build one hard-coded agent and print +// "All tests completed successfully", which claimed far more than it checked. +func runTests(out io.Writer, args []string) error { + llmManager := llm.NewProviderManager() + toolRegistry := tools.NewToolRegistry() - // Create test configuration - testConfig := &agent.AgentConfig{ - Name: "test-agent", - Type: agent.AgentTypeChat, - Model: "gpt-3.5-turbo", - Provider: "openai", - SystemPrompt: "You are a helpful assistant for testing.", - Temperature: 0.7, - MaxTokens: 100, + var configs []*agentFileConfig + if len(args) > 0 { + _, _ = fmt.Fprintf(out, "Testing agent configuration %s...\n", args[0]) + loaded, err := loadAgentConfigs(args[0]) + if err != nil { + return err + } + configs = loaded + } else { + _, _ = fmt.Fprintln(out, "No configuration given; running the built-in self-test...") + configs = []*agentFileConfig{{ + Name: "test-agent", + Type: string(agent.AgentTypeChat), + Model: "gpt-3.5-turbo", + Provider: "openai", + SystemPrompt: "You are a helpful assistant for testing.", + Temperature: 0.7, + MaxTokens: 1000, + }} } - // Initialize components for testing - llmManager := llm.NewProviderManager() - toolRegistry := tools.NewToolRegistry() + for _, cfg := range configs { + agentConfig := cfg.toAgentConfig() + // NewAgent falls back to defaults when handed an invalid config, so + // check the configuration itself before trusting the agent it returns. + if err := agentConfig.Validate(); err != nil { + return fmt.Errorf("agent %s: %w", cfg.Name, err) + } - // Create test agent - testAgent := agent.NewAgent(testConfig, llmManager, toolRegistry) + built := agent.NewAgent(agentConfig, llmManager, toolRegistry) + if got := built.GetConfig().Name; got != cfg.Name { + return fmt.Errorf("agent %s: configuration was rejected by the agent runtime (got name %q)", cfg.Name, got) + } - fmt.Printf("Test agent created: %s\n", testAgent.GetConfig().Name) - fmt.Printf("Agent type: %s\n", testAgent.GetConfig().Type) - fmt.Printf("Model: %s\n", testAgent.GetConfig().Model) + graph := built.GetGraph() + if graph == nil { + return fmt.Errorf("agent %s: no execution graph was built", cfg.Name) + } + if err := graph.Validate(); err != nil { + return fmt.Errorf("agent %s: graph validation failed: %w", cfg.Name, err) + } - // Validate graph structure - graph := testAgent.GetGraph() - if err := graph.Validate(); err != nil { - log.Fatalf("Graph validation failed: %v", err) + _, _ = fmt.Fprintf(out, " βœ“ %s (%s, %s/%s): graph valid\n", + cfg.Name, agentConfig.Type, agentConfig.Provider, agentConfig.Model) } - fmt.Println("Graph validation passed") - fmt.Println("All tests completed successfully") + _, _ = fmt.Fprintf(out, "%d agent configuration(s) built and validated. No LLM calls were made.\n", len(configs)) + return nil +} + +// safeProjectDir turns a project name into a directory below the working +// directory, rejecting names that escape it. +// +// "golanggraph init ../../etc/whatever" used to create and populate directories +// anywhere the process could write. +func safeProjectDir(name string) (string, error) { + if strings.TrimSpace(name) == "" { + return "", errors.New("project name must not be empty") + } + if filepath.IsAbs(name) { + return "", fmt.Errorf("project name %q must be relative to the current directory", name) + } + cleaned := filepath.Clean(name) + if cleaned == "." || cleaned == ".." || strings.HasPrefix(cleaned, ".."+string(filepath.Separator)) { + return "", fmt.Errorf("project name %q escapes the current directory", name) + } + return cleaned, nil } -func runInit(args []string, template string) { - fmt.Printf("Initializing new GoLangGraph project...\n") +// projectTemplates lists the templates init understands. +var projectTemplates = []string{"basic", "advanced", "rag"} +func runInit(out io.Writer, args []string, template string, force bool) error { projectName := "golanggraph-agent" if len(args) > 0 { projectName = args[0] } - fmt.Printf("Creating project: %s with template: %s\n", projectName, template) + dir, err := safeProjectDir(projectName) + if err != nil { + return err + } - // Create project directory - if err := os.MkdirAll(projectName, 0750); err != nil { - log.Fatalf("Failed to create project directory: %v", err) + // An unknown template used to fall through to the basic one while still + // reporting the template the operator asked for. + valid := false + for _, t := range projectTemplates { + if t == template { + valid = true + break + } + } + if !valid { + return fmt.Errorf("unknown template %q (want one of: %s)", template, strings.Join(projectTemplates, ", ")) } - // Create subdirectories - dirs := []string{ - "configs", - "agents", - "tools", - "static", - "tests", + if entries, readErr := os.ReadDir(dir); readErr == nil && len(entries) > 0 && !force { + return fmt.Errorf("directory %s already exists and is not empty (use --force to overwrite)", dir) } - for _, dir := range dirs { - if err := os.MkdirAll(fmt.Sprintf("%s/%s", projectName, dir), 0750); err != nil { - log.Fatalf("Failed to create directory %s: %v", dir, err) + _, _ = fmt.Fprintf(out, "Creating project %s from the %s template...\n", dir, template) + + for _, sub := range []string{"", "configs", "agents", "tools", "static", "tests"} { + if mkErr := os.MkdirAll(filepath.Join(dir, sub), 0750); mkErr != nil { + return fmt.Errorf("failed to create directory %s: %w", filepath.Join(dir, sub), mkErr) + } + } + + // Every project gets a buildable Go program: init used to produce a + // directory of YAML with no code in it and then tell the operator to run it. + files := []struct { + path string + content string + }{ + {"go.mod", projectGoMod(filepath.Base(dir))}, + {"main.go", projectMainGo(filepath.Base(dir))}, + {"README.md", projectReadme(filepath.Base(dir), template)}, + {".gitignore", "/" + filepath.Base(dir) + "\n*.exe\n.env\n"}, + } + for _, f := range files { + if writeErr := writeFileChecked(filepath.Join(dir, f.path), f.content); writeErr != nil { + return writeErr } } - // Create template files based on template type switch template { - case "basic": - createBasicTemplate(projectName) case "advanced": - createAdvancedTemplate(projectName) + err = createAdvancedTemplate(dir) case "rag": - createRAGTemplate(projectName) + err = createRAGTemplate(dir) default: - createBasicTemplate(projectName) + err = createBasicTemplate(dir) + } + if err != nil { + return err + } + + _, _ = fmt.Fprintf(out, "Project %s initialized.\n", dir) + _, _ = fmt.Fprintf(out, "Next steps:\n") + _, _ = fmt.Fprintf(out, " cd %s\n", dir) + _, _ = fmt.Fprintf(out, " go mod tidy && go run .\n") + _, _ = fmt.Fprintf(out, " golanggraph validate configs/agent-config.yaml\n") + return nil +} + +// writeFileChecked writes a generated project file and reports failure. Silent +// os.WriteFile errors are how scaffolding ends up claiming files it never made. +func writeFileChecked(path, content string) error { + if err := os.WriteFile(path, []byte(content), 0600); err != nil { + return fmt.Errorf("failed to write %s: %w", path, err) + } + return nil +} + +func projectGoMod(name string) string { + return fmt.Sprintf(`module %s + +go 1.23.0 + +require github.com/UnicoLab/GoLangGraph v0.0.0 + +// The framework is not published to a module proxy yet; point this at your +// checkout (or delete both lines once it is). +replace github.com/UnicoLab/GoLangGraph => ../GoLangGraph +`, name) +} + +func projectMainGo(name string) string { + return fmt.Sprintf(`// Command %s is a GoLangGraph agent generated by "golanggraph init". +package main + +import ( + "context" + "fmt" + "log" + "os" + "time" + + "github.com/UnicoLab/GoLangGraph/pkg/agent" + "github.com/UnicoLab/GoLangGraph/pkg/llm" + "github.com/UnicoLab/GoLangGraph/pkg/tools" +) + +func main() { + endpoint := os.Getenv("OLLAMA_URL") + if endpoint == "" { + endpoint = "http://localhost:11434" + } + + provider, err := llm.NewOllamaProvider(&llm.ProviderConfig{ + Endpoint: endpoint, + Timeout: 60 * time.Second, + }) + if err != nil { + log.Fatalf("failed to create the ollama provider: %%v", err) } - fmt.Printf("Project %s initialized successfully!\n", projectName) - fmt.Printf("Next steps:\n") - fmt.Printf(" cd %s\n", projectName) - fmt.Printf(" golanggraph dev\n") + providers := llm.NewProviderManager() + if err := providers.RegisterProvider("ollama", provider); err != nil { + log.Fatalf("failed to register the ollama provider: %%v", err) + } + + config := agent.DefaultAgentConfig() + config.Name = %q + config.Type = agent.AgentTypeChat + config.Provider = "ollama" + config.Model = "gemma3:1b" + config.SystemPrompt = "You are a helpful assistant." + + if err := config.Validate(); err != nil { + log.Fatalf("invalid agent configuration: %%v", err) + } + + assistant := agent.NewAgent(config, providers, tools.NewToolRegistry()) + + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute) + defer cancel() + + execution, err := assistant.Execute(ctx, "Say hello in one sentence.") + if err != nil { + log.Fatalf("agent execution failed: %%v", err) + } + + fmt.Println(execution.Output) +} +`, name, name) } -func runDockerBuild(args []string, distroless bool, tag, dockerfile, platform string) { - fmt.Printf("Building Docker container...\n") +func projectReadme(name, template string) string { + return fmt.Sprintf(`# %s + +A GoLangGraph project generated with the %q template. + +## Run + + go mod tidy + go run . + +The generated agent talks to a local Ollama (set OLLAMA_URL to point elsewhere). + +## Layout + +- main.go the agent program +- configs/ agent configuration files +- docker-compose.yml postgres and redis for state persistence + +## Validate the configuration + + golanggraph validate configs/agent-config.yaml +`, name, template) +} + +// dockerBuildOptions describes a docker build run. +type dockerBuildOptions struct { + Distroless bool + Tag string + Dockerfile string + Platform string + DryRun bool + ContextDir string +} + +// runCommand executes an external command. It is a package variable so tests +// can observe the command line without needing docker installed. +var runCommand = func(ctx context.Context, out io.Writer, name string, args ...string) error { + cmd := exec.CommandContext(ctx, name, args...) // #nosec G204 -- arguments are built from validated flags + cmd.Stdout = out + cmd.Stderr = out + return cmd.Run() +} + +// runDockerBuild builds the container image. +// +// It used to print the command it would have run and then report "Docker build +// command prepared" and exit 0, so a build pipeline calling it produced no +// image and no error. It now runs docker (use --dry-run to only print), and +// fails when docker is missing or the build fails. +func runDockerBuild(ctx context.Context, out io.Writer, args []string, opts dockerBuildOptions) error { + if ctx == nil { + ctx = context.Background() + } configFile := "agent-config.yaml" if len(args) > 0 { configFile = args[0] } + // Building an image around a configuration that does not parse just moves + // the failure to production. + if _, err := loadAgentConfigs(configFile); err != nil { + return fmt.Errorf("agent config %s: %w", configFile, err) + } - fmt.Printf("Using config file: %s\n", configFile) - - // Determine image tag + tag := opts.Tag if tag == "" { tag = "golanggraph-agent:latest" } + contextDir := opts.ContextDir + if contextDir == "" { + contextDir = "." + } - // Choose dockerfile based on distroless flag var dockerfilePath string - if dockerfile != "" { - dockerfilePath = dockerfile - } else if distroless { + switch { + case opts.Dockerfile != "": + dockerfilePath = opts.Dockerfile + if _, err := os.Stat(dockerfilePath); err != nil { + return fmt.Errorf("dockerfile %s: %w", dockerfilePath, err) + } + case opts.Distroless: dockerfilePath = "Dockerfile.distroless" - createDistrolessDockerfile(dockerfilePath) - } else { + if err := ensureDockerfile(out, dockerfilePath, distrolessDockerfile); err != nil { + return err + } + default: dockerfilePath = "Dockerfile.agent" - createAgentDockerfile(dockerfilePath) + if err := ensureDockerfile(out, dockerfilePath, agentDockerfile); err != nil { + return err + } + } + + dockerArgs := []string{"build", "-f", dockerfilePath, "-t", tag} + if opts.Platform != "" { + dockerArgs = append(dockerArgs, "--platform", opts.Platform) } + dockerArgs = append(dockerArgs, contextDir) - // Build Docker command - var dockerCmd []string - dockerCmd = append(dockerCmd, "docker", "build", "-f", dockerfilePath, "-t", tag) + _, _ = fmt.Fprintf(out, "docker %s\n", strings.Join(dockerArgs, " ")) - if platform != "" { - dockerCmd = append(dockerCmd, "--platform", platform) + if opts.DryRun { + _, _ = fmt.Fprintln(out, "Dry run: the image was not built.") + return nil } - dockerCmd = append(dockerCmd, ".") + if err := runCommand(ctx, out, "docker", dockerArgs...); err != nil { + return fmt.Errorf("docker build failed: %w", err) + } - fmt.Printf("Running: %s\n", fmt.Sprintf("%v", dockerCmd)) - fmt.Printf("Image tag: %s\n", tag) - fmt.Printf("Dockerfile: %s\n", dockerfilePath) - fmt.Printf("Distroless: %t\n", distroless) + _, _ = fmt.Fprintf(out, "Built image %s\n", tag) + return nil +} - // Note: In a real implementation, you would execute the docker command - // For now, we'll just show what would be executed - fmt.Printf("Docker build command prepared. Execute manually or integrate with docker library.\n") +// ensureDockerfile writes the generated Dockerfile, leaving an existing one +// alone: silently overwriting an operator's Dockerfile loses their changes. +func ensureDockerfile(out io.Writer, path, content string) error { + if _, err := os.Stat(path); err == nil { + _, _ = fmt.Fprintf(out, "Using existing %s\n", path) + return nil + } else if !errors.Is(err, os.ErrNotExist) { + return fmt.Errorf("cannot inspect %s: %w", path, err) + } + + if err := writeFileChecked(path, content); err != nil { + return err + } + _, _ = fmt.Fprintf(out, "Generated %s\n", path) + return nil } -func runDevServer() { - fmt.Println("Starting development server...") +// agentFileConfig is one agent as it appears in a configuration file. The +// framework's AgentConfig carries JSON tags only, so decoding YAML straight +// into it drops every snake_case key ("system_prompt", "max_tokens") without a +// word of complaint. This type accepts both spellings. +type agentFileConfig struct { + Key string // map key in a multi-agent file, empty for single-agent files + ID string + Name string + Type string + Model string + Provider string + SystemPrompt string + Temperature float64 + MaxTokens int + MaxIterations int + Tools []string + UnknownKeys []string +} - // Create development server configuration - config := &server.ServerConfig{ - Host: viper.GetString("host"), - Port: viper.GetInt("port"), - ReadTimeout: 30 * time.Second, - WriteTimeout: 30 * time.Second, - MaxHeaderBytes: 1 << 20, - EnableCORS: true, - StaticDir: "./static", - DevMode: true, +// toAgentConfig converts to the framework configuration, filling in the +// framework defaults for anything the file left out. +func (c *agentFileConfig) toAgentConfig() *agent.AgentConfig { + config := agent.DefaultAgentConfig() + if c.ID != "" { + config.ID = c.ID + } + config.Name = c.Name + if c.Type != "" { + config.Type = agent.AgentType(c.Type) + } + config.Model = c.Model + config.Provider = c.Provider + config.SystemPrompt = c.SystemPrompt + if c.Temperature != 0 { + config.Temperature = c.Temperature + } + if c.MaxTokens != 0 { + config.MaxTokens = c.MaxTokens + } + if c.MaxIterations != 0 { + config.MaxIterations = c.MaxIterations } + if c.Tools != nil { + config.Tools = c.Tools + } + return config +} - // Create server - srv := server.NewServer(config) +// knownAgentKeys are the keys understood inside an agent block. +var knownAgentKeys = map[string]bool{ + "id": true, "name": true, "type": true, "model": true, "provider": true, + "systemprompt": true, "temperature": true, "maxtokens": true, + "maxiterations": true, "tools": true, "enablestreaming": true, + "streamingmode": true, "timeout": true, "metadata": true, + "description": true, "enabled": true, +} + +// knownTopLevelKeys are the keys understood beside the agent definition in a +// configuration file. +var knownTopLevelKeys = map[string]bool{ + "agents": true, "routing": true, "deployment": true, "shared": true, + "version": true, "database": true, "vectorstore": true, "rag": true, + "documentloaders": true, "workflow": true, "server": true, +} - // Initialize components - if err := initializeComponents(srv); err != nil { - log.Fatalf("Failed to initialize components: %v", err) +// normalizeKey folds "system_prompt", "system-prompt" and "systemPrompt" onto +// one spelling so a configuration is not silently half-read. +func normalizeKey(key string) string { + return strings.ToLower(strings.NewReplacer("_", "", "-", "", " ", "").Replace(key)) +} + +// decodeConfigFile reads a YAML or JSON configuration file into a generic map. +func decodeConfigFile(path string) (map[string]interface{}, error) { + data, err := os.ReadFile(path) // #nosec G304 -- reading the operator's configuration file is the command + if err != nil { + return nil, err + } + if len(strings.TrimSpace(string(data))) == 0 { + return nil, fmt.Errorf("%s is empty", path) } - // Start server in a goroutine - go func() { - if err := srv.Start(); err != nil { - log.Fatalf("Server failed to start: %v", err) + raw := map[string]interface{}{} + switch ext := strings.ToLower(filepath.Ext(path)); ext { + case ".json": + if err := json.Unmarshal(data, &raw); err != nil { + return nil, fmt.Errorf("parsing %s: %w", path, err) + } + case ".yaml", ".yml": + if err := yaml.Unmarshal(data, &raw); err != nil { + return nil, fmt.Errorf("parsing %s: %w", path, err) } - }() + default: + return nil, fmt.Errorf("unsupported config extension %q (want .yaml, .yml or .json)", ext) + } + if raw == nil { + return nil, fmt.Errorf("%s does not contain a configuration object", path) + } + return raw, nil +} - fmt.Printf("Development server started on %s:%d\n", config.Host, config.Port) - fmt.Printf("API endpoints: http://%s:%d/api/v1/\n", config.Host, config.Port) - fmt.Printf("Debug interface: http://%s:%d/debug\n", config.Host, config.Port) - fmt.Printf("Agent playground: http://%s:%d/playground\n", config.Host, config.Port) +// parseAgentBlock reads one agent definition out of a decoded map. +func parseAgentBlock(key string, raw map[string]interface{}) (*agentFileConfig, error) { + cfg := &agentFileConfig{Key: key} + + for rawKey, value := range raw { + switch name := normalizeKey(rawKey); name { + case "id": + cfg.ID = fmt.Sprintf("%v", value) + case "name": + cfg.Name = fmt.Sprintf("%v", value) + case "type": + cfg.Type = fmt.Sprintf("%v", value) + case "model": + cfg.Model = fmt.Sprintf("%v", value) + case "provider": + cfg.Provider = fmt.Sprintf("%v", value) + case "systemprompt": + cfg.SystemPrompt = fmt.Sprintf("%v", value) + case "temperature": + f, err := toFloat(value) + if err != nil { + return nil, fmt.Errorf("temperature: %w", err) + } + cfg.Temperature = f + case "maxtokens": + n, err := toInt(value) + if err != nil { + return nil, fmt.Errorf("max_tokens: %w", err) + } + cfg.MaxTokens = n + case "maxiterations": + n, err := toInt(value) + if err != nil { + return nil, fmt.Errorf("max_iterations: %w", err) + } + cfg.MaxIterations = n + case "tools": + names, err := toToolNames(value) + if err != nil { + return nil, fmt.Errorf("tools: %w", err) + } + cfg.Tools = names + default: + if !knownAgentKeys[name] && !knownTopLevelKeys[name] { + cfg.UnknownKeys = append(cfg.UnknownKeys, rawKey) + } + } + } - // Watch for file changes (hot-reload) - if viper.GetBool("hot-reload") { - fmt.Println("Hot-reload enabled - watching for changes...") - // Note: File watching implementation would go here + sort.Strings(cfg.UnknownKeys) + if cfg.Name == "" { + cfg.Name = key } + return cfg, nil +} - // Wait for interrupt signal to gracefully shutdown - quit := make(chan os.Signal, 1) - signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM) - <-quit +func toFloat(value interface{}) (float64, error) { + switch v := value.(type) { + case float64: + return v, nil + case float32: + return float64(v), nil + case int: + return float64(v), nil + case int64: + return float64(v), nil + default: + return 0, fmt.Errorf("want a number, got %T (%v)", value, value) + } +} - fmt.Println("Shutting down development server...") +func toInt(value interface{}) (int, error) { + switch v := value.(type) { + case int: + return v, nil + case int64: + return int(v), nil + case float64: + if v != float64(int(v)) { + return 0, fmt.Errorf("want a whole number, got %v", v) + } + return int(v), nil + default: + return 0, fmt.Errorf("want a whole number, got %T (%v)", value, value) + } +} - // Create a deadline for shutdown - ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) - defer cancel() +// toToolNames accepts both `tools: [calculator]` and the list-of-objects form +// `tools: [{name: calculator, enabled: true}]` that the init template writes. +func toToolNames(value interface{}) ([]string, error) { + items, ok := value.([]interface{}) + if !ok { + return nil, fmt.Errorf("want a list, got %T", value) + } + + var names []string + for _, item := range items { + switch entry := item.(type) { + case string: + names = append(names, entry) + case map[string]interface{}: + enabled := true + var name string + for k, v := range entry { + switch normalizeKey(k) { + case "name": + name = fmt.Sprintf("%v", v) + case "enabled": + if b, ok := v.(bool); ok { + enabled = b + } + } + } + if name == "" { + return nil, fmt.Errorf("tool entry %v has no name", entry) + } + if enabled { + names = append(names, name) + } + default: + return nil, fmt.Errorf("want a tool name or {name, enabled}, got %T", item) + } + } + return names, nil +} - if err := srv.Stop(ctx); err != nil { - log.Fatalf("Server forced to shutdown: %v", err) +// loadAgentConfigs parses every agent defined in a configuration file. Both a +// single-agent file and a multi-agent file (top-level "agents:") are accepted. +func loadAgentConfigs(path string) ([]*agentFileConfig, error) { + raw, err := decodeConfigFile(path) + if err != nil { + return nil, err } - fmt.Println("Development server stopped") + for key, value := range raw { + if normalizeKey(key) != "agents" { + continue + } + agentsMap, ok := value.(map[string]interface{}) + if !ok { + return nil, fmt.Errorf("%s: \"agents\" must be a map of agent id to agent definition", path) + } + if len(agentsMap) == 0 { + return nil, fmt.Errorf("%s: \"agents\" is empty", path) + } + + ids := make([]string, 0, len(agentsMap)) + for id := range agentsMap { + ids = append(ids, id) + } + sort.Strings(ids) + + configs := make([]*agentFileConfig, 0, len(ids)) + for _, id := range ids { + block, ok := agentsMap[id].(map[string]interface{}) + if !ok { + return nil, fmt.Errorf("%s: agent %q is not a mapping", path, id) + } + cfg, parseErr := parseAgentBlock(id, block) + if parseErr != nil { + return nil, fmt.Errorf("%s: agent %q: %w", path, id, parseErr) + } + configs = append(configs, cfg) + } + return configs, nil + } + + cfg, err := parseAgentBlock("", raw) + if err != nil { + return nil, fmt.Errorf("%s: %w", path, err) + } + return []*agentFileConfig{cfg}, nil +} + +// validationReport collects everything found in a configuration. +type validationReport struct { + Errors []string + Warnings []string } -func runValidate(args []string, strict bool) { - fmt.Printf("Validating configuration...\n") +func (r *validationReport) errorf(format string, args ...interface{}) { + r.Errors = append(r.Errors, fmt.Sprintf(format, args...)) +} + +func (r *validationReport) warnf(format string, args ...interface{}) { + r.Warnings = append(r.Warnings, fmt.Sprintf(format, args...)) +} + +// knownProviders are the providers the framework ships with. +var knownProviders = map[string]bool{"openai": true, "ollama": true, "gemini": true} + +// validateAgentConfigs checks parsed agents for real problems: required fields, +// value ranges, known agent types, resolvable tools, and a graph that builds. +// +// The validate command used to check only that the file existed and then print +// "Configuration validation completed successfully!" -- it reported unparseable +// YAML as valid. +func validateAgentConfigs(configs []*agentFileConfig) *validationReport { + report := &validationReport{} + toolRegistry := tools.NewToolRegistry() + known := map[string]bool{} + for _, name := range toolRegistry.ListTools() { + known[name] = true + } + llmManager := llm.NewProviderManager() + + for _, cfg := range configs { + label := cfg.Name + if cfg.Key != "" { + label = cfg.Key + } + if label == "" { + label = "agent" + } + + agentConfig := cfg.toAgentConfig() + if err := agentConfig.Validate(); err != nil { + report.errorf("%s: %v", label, err) + continue + } + + // An unrecognized type is not rejected by the framework: it silently + // builds a chat agent, so "type: reactt" would run as something else. + switch agentConfig.Type { + case agent.AgentTypeChat, agent.AgentTypeReAct, agent.AgentTypeTool: + default: + report.errorf("%s: unknown agent type %q (want chat, react or tool)", label, agentConfig.Type) + continue + } + if !knownProviders[strings.ToLower(agentConfig.Provider)] { + report.warnf("%s: provider %q is not one of the built-in providers", label, agentConfig.Provider) + } + if agentConfig.SystemPrompt == "" { + report.warnf("%s: no system prompt", label) + } + for _, tool := range agentConfig.Tools { + if !known[tool] { + report.warnf("%s: tool %q is not registered", label, tool) + } + } + for _, key := range cfg.UnknownKeys { + report.warnf("%s: unknown key %q", label, key) + } + + built := agent.NewAgent(agentConfig, llmManager, toolRegistry) + graph := built.GetGraph() + if graph == nil { + report.errorf("%s: no execution graph was built", label) + continue + } + if err := graph.Validate(); err != nil { + report.errorf("%s: execution graph is invalid: %v", label, err) + } + } + + return report +} + +func runValidate(out io.Writer, args []string, strict bool) error { configFile := "agent-config.yaml" if len(args) > 0 { configFile = args[0] } - fmt.Printf("Config file: %s\n", configFile) - fmt.Printf("Strict mode: %t\n", strict) + _, _ = fmt.Fprintf(out, "Validating %s (strict: %t)...\n", configFile, strict) - // Check if config file exists - if _, err := os.Stat(configFile); os.IsNotExist(err) { - log.Fatalf("Configuration file not found: %s", configFile) + configs, err := loadAgentConfigs(configFile) + if err != nil { + return err } - // Note: In a real implementation, you would: - // 1. Parse the configuration file - // 2. Validate the schema - // 3. Check for required fields - // 4. Validate graph structure - // 5. Check tool availability - // 6. Validate LLM provider configuration + report := validateAgentConfigs(configs) - fmt.Printf("Configuration validation completed successfully!\n") + // Routing rules that point at agents which do not exist are a deployment + // outage waiting to happen, so check them here too. + if raw, err := decodeConfigFile(configFile); err == nil { + checkRouting(raw, configs, report) + } + + for _, warning := range report.Warnings { + _, _ = fmt.Fprintf(out, " ⚠ %s\n", warning) + } + for _, problem := range report.Errors { + _, _ = fmt.Fprintf(out, " βœ— %s\n", problem) + } + + if len(report.Errors) > 0 { + return fmt.Errorf("%s is invalid: %d problem(s)", configFile, len(report.Errors)) + } + if strict && len(report.Warnings) > 0 { + return fmt.Errorf("%s has %d warning(s) and --strict is set", configFile, len(report.Warnings)) + } + + _, _ = fmt.Fprintf(out, "βœ… %s is valid: %d agent(s), %d warning(s)\n", configFile, len(configs), len(report.Warnings)) + return nil } -func runDeployDocker(args []string) { - fmt.Printf("Deploying agent using Docker...\n") +// checkRouting verifies that routing rules reference agents that exist. +func checkRouting(raw map[string]interface{}, configs []*agentFileConfig, report *validationReport) { + ids := map[string]bool{} + for _, cfg := range configs { + if cfg.Key != "" { + ids[cfg.Key] = true + } + if cfg.ID != "" { + ids[cfg.ID] = true + } + } + if len(ids) == 0 { + return + } + + routing, ok := raw["routing"].(map[string]interface{}) + if !ok { + return + } + + if def, isString := routing["default_agent"].(string); isString && def != "" && !ids[def] { + report.errorf("routing: default agent %q is not defined", def) + } + rules, ok := routing["rules"].([]interface{}) + if !ok { + return + } + patterns := map[string]string{} + for _, item := range rules { + rule, ok := item.(map[string]interface{}) + if !ok { + continue + } + agentID, _ := rule["agent_id"].(string) + if agentID != "" && !ids[agentID] { + report.errorf("routing: rule targets agent %q, which is not defined", agentID) + } + if pattern, ok := rule["pattern"].(string); ok && pattern != "" { + if previous, clash := patterns[pattern]; clash { + report.warnf("routing: pattern %q is used by both %q and %q", pattern, previous, agentID) + } + patterns[pattern] = agentID + } + } +} + +// runDeployDocker refuses to claim a deployment it cannot perform. +// +// This command used to print "Docker deployment completed for config: X!" for +// any argument at all -- including a path that does not exist -- while doing +// nothing whatsoever. +func runDeployDocker(out io.Writer, args []string) error { configFile := "agent-config.yaml" if len(args) > 0 { configFile = args[0] } - fmt.Printf("Config file: %s\n", configFile) + configs, err := loadAgentConfigs(configFile) + if err != nil { + return fmt.Errorf("agent config %s: %w", configFile, err) + } - // Note: In a real implementation, you would: - // 1. Build the Docker image - // 2. Push to registry - // 3. Deploy to target environment - // 4. Monitor deployment status + report := validateAgentConfigs(configs) + if len(report.Errors) > 0 { + for _, problem := range report.Errors { + _, _ = fmt.Fprintf(out, " βœ— %s\n", problem) + } + return fmt.Errorf("%s is invalid: %d problem(s)", configFile, len(report.Errors)) + } - fmt.Printf("Docker deployment completed for config: %s!\n", configFile) + _, _ = fmt.Fprintf(out, "%s is valid (%d agent(s)).\n", configFile, len(configs)) + return fmt.Errorf("deploying to docker is %w: build an image with 'golanggraph docker build %s' and run it with docker or docker compose", errNotImplemented, configFile) } - -func createBasicTemplate(projectName string) { +func createBasicTemplate(projectName string) error { // Create basic agent configuration agentConfig := `name: "basic-agent" type: "chat" @@ -780,8 +1791,8 @@ database: password: "password" ` - if err := os.WriteFile(fmt.Sprintf("%s/configs/agent-config.yaml", projectName), []byte(agentConfig), 0600); err != nil { - log.Fatalf("Failed to create agent config: %v", err) + if err := writeFileChecked(filepath.Join(projectName, "configs", "agent-config.yaml"), agentConfig); err != nil { + return err } // Create docker-compose for development @@ -807,13 +1818,13 @@ volumes: postgres_data: ` - if err := os.WriteFile(fmt.Sprintf("%s/docker-compose.yml", projectName), []byte(dockerCompose), 0600); err != nil { - log.Fatalf("Failed to create docker-compose: %v", err) - } + return writeFileChecked(filepath.Join(projectName, "docker-compose.yml"), dockerCompose) } -func createAdvancedTemplate(projectName string) { - createBasicTemplate(projectName) +func createAdvancedTemplate(projectName string) error { + if err := createBasicTemplate(projectName); err != nil { + return err + } // Add advanced configuration advancedConfig := `name: "advanced-agent" @@ -878,13 +1889,13 @@ vector_store: dimensions: 1536 ` - if err := os.WriteFile(fmt.Sprintf("%s/configs/advanced-config.yaml", projectName), []byte(advancedConfig), 0600); err != nil { - log.Fatalf("Failed to create advanced config: %v", err) - } + return writeFileChecked(filepath.Join(projectName, "configs", "advanced-config.yaml"), advancedConfig) } -func createRAGTemplate(projectName string) { - createAdvancedTemplate(projectName) +func createRAGTemplate(projectName string) error { + if err := createAdvancedTemplate(projectName); err != nil { + return err + } // Add RAG-specific configuration ragConfig := `name: "rag-agent" @@ -942,13 +1953,11 @@ database: password: "password" ` - if err := os.WriteFile(fmt.Sprintf("%s/configs/rag-config.yaml", projectName), []byte(ragConfig), 0600); err != nil { - log.Fatalf("Failed to create RAG config: %v", err) - } + return writeFileChecked(filepath.Join(projectName, "configs", "rag-config.yaml"), ragConfig) } -func createAgentDockerfile(filepath string) { - dockerfile := `# Production Dockerfile for GoLangGraph Agent +// agentDockerfile is the Dockerfile generated for a normal build. +const agentDockerfile = `# Production Dockerfile for GoLangGraph Agent FROM golang:1.21-alpine AS builder # Set working directory @@ -1014,13 +2023,8 @@ ENTRYPOINT ["./golanggraph-agent"] CMD ["serve", "--host", "0.0.0.0", "--port", "8080"] ` - if err := os.WriteFile(filepath, []byte(dockerfile), 0600); err != nil { - log.Fatalf("Failed to create Dockerfile: %v", err) - } -} - -func createDistrolessDockerfile(filepath string) { - dockerfile := `# Distroless Dockerfile for GoLangGraph Agent +// distrolessDockerfile is the Dockerfile generated for --distroless builds. +const distrolessDockerfile = `# Distroless Dockerfile for GoLangGraph Agent FROM golang:1.21-alpine AS builder # Set working directory @@ -1070,77 +2074,9 @@ ENTRYPOINT ["/golanggraph-agent"] CMD ["serve", "--host", "0.0.0.0", "--port", "8080"] ` - if err := os.WriteFile(filepath, []byte(dockerfile), 0600); err != nil { - log.Fatalf("Failed to create distroless Dockerfile: %v", err) - } -} - -func runHealthCheck() { - fmt.Printf("Running GoLangGraph health check...\n") - - healthy := true - var issues []string - - // Check system resources - fmt.Printf("Checking system resources...\n") - - // Check database connectivity - fmt.Printf("Checking database connectivity...\n") - dbHost := os.Getenv("POSTGRES_HOST") - if dbHost == "" { - dbHost = "localhost" - } - fmt.Printf(" PostgreSQL: %s:5432 - ", dbHost) - // In a real implementation, you would test actual connectivity - fmt.Printf("βœ“ Reachable\n") - - redisHost := os.Getenv("REDIS_HOST") - if redisHost == "" { - redisHost = "localhost" - } - fmt.Printf(" Redis: %s:6379 - ", redisHost) - // In a real implementation, you would test actual connectivity - fmt.Printf("βœ“ Reachable\n") - - // Check LLM providers - fmt.Printf("Checking LLM providers...\n") - if apiKey := os.Getenv("OPENAI_API_KEY"); apiKey != "" { - fmt.Printf(" OpenAI: βœ“ API key configured\n") - } else { - fmt.Printf(" OpenAI: ⚠ API key not configured\n") - issues = append(issues, "OpenAI API key not configured") - } - - ollamaURL := os.Getenv("OLLAMA_URL") - if ollamaURL == "" { - ollamaURL = "http://localhost:11434" - } - fmt.Printf(" Ollama: %s - ", ollamaURL) - // In a real implementation, you would test actual connectivity - fmt.Printf("βœ“ Reachable\n") - - // Check disk space - fmt.Printf("Checking system resources...\n") - fmt.Printf(" Disk space: βœ“ Sufficient\n") - fmt.Printf(" Memory: βœ“ Available\n") - - // Overall health status - fmt.Printf("\n") - if healthy && len(issues) == 0 { - fmt.Printf("βœ… System is healthy\n") - os.Exit(0) - } else { - fmt.Printf("⚠ System has issues:\n") - for _, issue := range issues { - fmt.Printf(" - %s\n", issue) - } - os.Exit(1) - } -} - func main() { if err := rootCmd.Execute(); err != nil { - fmt.Fprintf(os.Stderr, "Error: %v\n", err) + _, _ = fmt.Fprintf(os.Stderr, "Error: %v\n", err) os.Exit(1) } } diff --git a/cmd/golanggraph/multi_agent_commands.go b/cmd/golanggraph/multi_agent_commands.go index 1bf8636..4206477 100644 --- a/cmd/golanggraph/multi_agent_commands.go +++ b/cmd/golanggraph/multi_agent_commands.go @@ -10,13 +10,16 @@ import ( "context" "encoding/json" "fmt" + "io" "os" + "os/signal" "path/filepath" + "sort" "strings" + "syscall" "time" "github.com/spf13/cobra" - "github.com/spf13/viper" yaml "gopkg.in/yaml.v3" "github.com/UnicoLab/GoLangGraph/pkg/agent" @@ -48,12 +51,24 @@ var multiAgentInitCmd = &cobra.Command{ Creates a directory structure optimized for managing multiple agents with different configurations, routing rules, and deployment settings.`, Args: cobra.MaximumNArgs(1), - Run: func(cmd *cobra.Command, args []string) { - template, _ := cmd.Flags().GetString("template") - agentCount, _ := cmd.Flags().GetInt("agents") - outputFormat, _ := cmd.Flags().GetString("format") - routingType, _ := cmd.Flags().GetString("routing") - runMultiAgentInit(args, template, agentCount, outputFormat, routingType) + RunE: func(cmd *cobra.Command, args []string) error { + template, err := cmd.Flags().GetString("template") + if err != nil { + return err + } + agentCount, err := cmd.Flags().GetInt("agents") + if err != nil { + return err + } + outputFormat, err := cmd.Flags().GetString("format") + if err != nil { + return err + } + routingType, err := cmd.Flags().GetString("routing") + if err != nil { + return err + } + return runMultiAgentInit(cmd.OutOrStdout(), args, template, agentCount, outputFormat, routingType) }, } @@ -64,10 +79,16 @@ var multiAgentValidateCmd = &cobra.Command{ Long: `Validate multi-agent configuration files including agent definitions, routing rules, deployment settings, and schema validation.`, Args: cobra.MaximumNArgs(1), - Run: func(cmd *cobra.Command, args []string) { - strict, _ := cmd.Flags().GetBool("strict") - checkSchemas, _ := cmd.Flags().GetBool("check-schemas") - runMultiAgentValidate(args, strict, checkSchemas) + RunE: func(cmd *cobra.Command, args []string) error { + strict, err := cmd.Flags().GetBool("strict") + if err != nil { + return err + } + checkSchemas, err := cmd.Flags().GetBool("check-schemas") + if err != nil { + return err + } + return runMultiAgentValidate(cmd.OutOrStdout(), args, strict, checkSchemas) }, } @@ -78,12 +99,20 @@ var multiAgentDeployCmd = &cobra.Command{ Long: `Deploy multiple agents according to the multi-agent configuration. Supports various deployment targets including Docker, Kubernetes, and serverless platforms.`, Args: cobra.MaximumNArgs(1), - Run: func(cmd *cobra.Command, args []string) { - deploymentType, _ := cmd.Flags().GetString("type") - environment, _ := cmd.Flags().GetString("environment") - dryRun, _ := cmd.Flags().GetBool("dry-run") - parallel, _ := cmd.Flags().GetBool("parallel") - runMultiAgentDeploy(args, deploymentType, environment, dryRun, parallel) + RunE: func(cmd *cobra.Command, args []string) error { + deploymentType, err := cmd.Flags().GetString("type") + if err != nil { + return err + } + environment, err := cmd.Flags().GetString("environment") + if err != nil { + return err + } + dryRun, err := cmd.Flags().GetBool("dry-run") + if err != nil { + return err + } + return runMultiAgentDeploy(cmd.OutOrStdout(), args, deploymentType, environment, dryRun) }, } @@ -94,10 +123,18 @@ var multiAgentServeCmd = &cobra.Command{ Long: `Start a server that hosts multiple agents with routing and load balancing. Provides HTTP endpoints for agent execution, management, and monitoring.`, Args: cobra.MaximumNArgs(1), - Run: func(cmd *cobra.Command, args []string) { - host, _ := cmd.Flags().GetString("host") - port, _ := cmd.Flags().GetInt("port") - runMultiAgentServe(args, host, port) + RunE: func(cmd *cobra.Command, args []string) error { + host, err := cmd.Flags().GetString("host") + if err != nil { + return err + } + port, err := cmd.Flags().GetInt("port") + if err != nil { + return err + } + ctx, stop := signal.NotifyContext(cmd.Context(), syscall.SIGINT, syscall.SIGTERM) + defer stop() + return runMultiAgentServe(ctx, cmd.OutOrStdout(), args, host, port) }, } @@ -107,10 +144,16 @@ var multiAgentStatusCmd = &cobra.Command{ Short: "Check status of deployed agents", Long: `Check the status of deployed agents including health, metrics, and runtime information.`, Args: cobra.MaximumNArgs(1), - Run: func(cmd *cobra.Command, args []string) { - outputFormat, _ := cmd.Flags().GetString("format") - watch, _ := cmd.Flags().GetBool("watch") - runMultiAgentStatus(args, outputFormat, watch) + RunE: func(cmd *cobra.Command, args []string) error { + outputFormat, err := cmd.Flags().GetString("format") + if err != nil { + return err + } + watch, err := cmd.Flags().GetBool("watch") + if err != nil { + return err + } + return runMultiAgentStatus(cmd.OutOrStdout(), args, outputFormat, watch) }, } @@ -127,10 +170,16 @@ var multiAgentGenerateDockerCmd = &cobra.Command{ Short: "Generate Docker deployment files", Long: `Generate Docker Compose files and Dockerfiles for multi-agent deployment.`, Args: cobra.MaximumNArgs(1), - Run: func(cmd *cobra.Command, args []string) { - outputDir, _ := cmd.Flags().GetString("output") - multiService, _ := cmd.Flags().GetBool("multi-service") - runGenerateDocker(args, outputDir, multiService) + RunE: func(cmd *cobra.Command, args []string) error { + outputDir, err := cmd.Flags().GetString("output") + if err != nil { + return err + } + multiService, err := cmd.Flags().GetBool("multi-service") + if err != nil { + return err + } + return runGenerateDocker(cmd.OutOrStdout(), args, outputDir, multiService) }, } @@ -140,10 +189,16 @@ var multiAgentGenerateK8sCmd = &cobra.Command{ Short: "Generate Kubernetes deployment manifests", Long: `Generate Kubernetes deployment, service, and ingress manifests for multi-agent deployment.`, Args: cobra.MaximumNArgs(1), - Run: func(cmd *cobra.Command, args []string) { - outputDir, _ := cmd.Flags().GetString("output") - namespace, _ := cmd.Flags().GetString("namespace") - runGenerateK8s(args, outputDir, namespace) + RunE: func(cmd *cobra.Command, args []string) error { + outputDir, err := cmd.Flags().GetString("output") + if err != nil { + return err + } + namespace, err := cmd.Flags().GetString("namespace") + if err != nil { + return err + } + return runGenerateK8s(cmd.OutOrStdout(), args, outputDir, namespace) }, } @@ -174,9 +229,11 @@ Examples: RunE: runMultiAgentLoad, } - multiAgentLoadCmd.Flags().BoolP("recursive", "r", false, "Recursively scan directories for Go files") - multiAgentLoadCmd.Flags().StringSliceP("include", "i", []string{"*.go"}, "File patterns to include") - multiAgentLoadCmd.Flags().StringSliceP("exclude", "e", []string{"*_test.go"}, "File patterns to exclude") + // These three configure directory scanning, which is not implemented; the + // command reports that rather than pretending to have loaded anything. + multiAgentLoadCmd.Flags().BoolP("recursive", "r", false, "Recursively scan directories (directory loading is not implemented)") + multiAgentLoadCmd.Flags().StringSliceP("include", "i", []string{"*.go"}, "File patterns to include (directory loading is not implemented)") + multiAgentLoadCmd.Flags().StringSliceP("exclude", "e", []string{"*_test.go"}, "File patterns to exclude (directory loading is not implemented)") multiAgentLoadCmd.Flags().BoolP("validate", "v", true, "Validate loaded agent definitions") multiAgentLoadCmd.Flags().BoolP("verbose", "", false, "Verbose output") @@ -227,7 +284,6 @@ This shows the source, type, and metadata for each registered agent.`, multiAgentDeployCmd.Flags().StringP("type", "t", "docker", "Deployment type (docker, kubernetes, serverless)") multiAgentDeployCmd.Flags().StringP("environment", "e", "development", "Deployment environment") multiAgentDeployCmd.Flags().Bool("dry-run", false, "Show what would be deployed without actually deploying") - multiAgentDeployCmd.Flags().Bool("parallel", true, "Deploy agents in parallel") // Multi-agent serve flags multiAgentServeCmd.Flags().StringP("host", "H", "0.0.0.0", "Host to bind to") @@ -247,166 +303,179 @@ This shows the source, type, and metadata for each registered agent.`, } // runMultiAgentInit initializes a new multi-agent project -func runMultiAgentInit(args []string, template string, agentCount int, outputFormat, routingType string) { +func runMultiAgentInit(out io.Writer, args []string, template string, agentCount int, outputFormat, routingType string) error { projectName := "golanggraph-multi-agent" if len(args) > 0 { projectName = args[0] } - fmt.Printf("Initializing multi-agent project: %s\n", projectName) - fmt.Printf("Template: %s, Agents: %d, Format: %s, Routing: %s\n", template, agentCount, outputFormat, routingType) - - // Create project directory - if err := os.MkdirAll(projectName, 0750); err != nil { - fmt.Printf("Error creating project directory: %v\n", err) - os.Exit(1) + // "multi-agent init ../../somewhere" used to scaffold outside the working + // directory; keep the project below it. + dir, err := safeProjectDir(projectName) + if err != nil { + return err } - // Create subdirectories - dirs := []string{ - "agents", - "configs", - "deploy", - "k8s", - "scripts", - "static", - "tests", + switch outputFormat { + case "yaml", "yml", "json": + default: + return fmt.Errorf("unsupported output format %q (want yaml or json)", outputFormat) } - - for _, dir := range dirs { - if err := os.MkdirAll(filepath.Join(projectName, dir), 0750); err != nil { - fmt.Printf("Error creating directory %s: %v\n", dir, err) - os.Exit(1) - } + switch template { + case "basic", "microservices", "rag", "workflow": + default: + return fmt.Errorf("unknown template %q (want basic, microservices, rag or workflow)", template) + } + switch routingType { + case "path", "host", "header", "query": + default: + return fmt.Errorf("unknown routing type %q (want path, host, header or query)", routingType) } + if agentCount <= 0 { + return fmt.Errorf("--agents must be at least 1, got %d", agentCount) + } + + _, _ = fmt.Fprintf(out, "Initializing multi-agent project: %s\n", dir) + _, _ = fmt.Fprintf(out, "Template: %s, Agents: %d, Format: %s, Routing: %s\n", template, agentCount, outputFormat, routingType) - // Create agent subdirectories - for i := 1; i <= agentCount; i++ { - agentDir := filepath.Join(projectName, "agents", fmt.Sprintf("agent-%d", i)) - if err := os.MkdirAll(agentDir, 0750); err != nil { - fmt.Printf("Error creating agent directory: %v\n", err) - os.Exit(1) + for _, sub := range []string{"", "agents", "configs", "deploy", "k8s", "scripts", "static", "tests"} { + if mkErr := os.MkdirAll(filepath.Join(dir, sub), 0750); mkErr != nil { + return fmt.Errorf("failed to create directory %s: %w", filepath.Join(dir, sub), mkErr) } } - // Generate multi-agent configuration config := createMultiAgentConfig(template, agentCount, routingType) - // Write configuration file - configFile := fmt.Sprintf("multi-agent.%s", outputFormat) - configPath := filepath.Join(projectName, "configs", configFile) - - var configData []byte - var err error - - switch outputFormat { - case "yaml", "yml": - configData, err = yaml.Marshal(config) - case "json": - configData, err = json.MarshalIndent(config, "", " ") - default: - fmt.Printf("Unsupported output format: %s\n", outputFormat) - os.Exit(1) + for agentID := range config.Agents { + if mkErr := os.MkdirAll(filepath.Join(dir, "agents", agentID), 0750); mkErr != nil { + return fmt.Errorf("failed to create agent directory: %w", mkErr) + } } + configFile := fmt.Sprintf("multi-agent.%s", outputFormat) + configData, err := marshalConfig(config, outputFormat) if err != nil { - fmt.Printf("Error marshaling configuration: %v\n", err) - os.Exit(1) + return fmt.Errorf("failed to encode configuration: %w", err) } - - if err := os.WriteFile(configPath, configData, 0600); err != nil { - fmt.Printf("Error writing configuration file: %v\n", err) - os.Exit(1) + if err := writeFileChecked(filepath.Join(dir, "configs", configFile), string(configData)); err != nil { + return err } - // Create individual agent configurations - createIndividualAgentConfigs(projectName, config, outputFormat) - - // Create Docker compose file - createDockerComposeFile(projectName, config) - - // Create Kubernetes manifests - createK8sManifests(projectName, config) + if err := createIndividualAgentConfigs(dir, config, outputFormat); err != nil { + return err + } + if err := createDockerComposeFile(dir, config); err != nil { + return err + } + if err := createK8sManifests(dir, config); err != nil { + return err + } + if err := createProjectREADME(dir, config); err != nil { + return err + } - // Create README - createProjectREADME(projectName, config) + _, _ = fmt.Fprintf(out, "\nMulti-agent project '%s' initialized.\n", dir) + _, _ = fmt.Fprintf(out, "\nNext steps:\n") + _, _ = fmt.Fprintf(out, " cd %s\n", dir) + _, _ = fmt.Fprintf(out, " golanggraph multi-agent validate configs/%s\n", configFile) + _, _ = fmt.Fprintf(out, " golanggraph multi-agent serve configs/%s\n", configFile) + return nil +} - fmt.Printf("\nMulti-agent project '%s' initialized successfully!\n", projectName) - fmt.Printf("\nNext steps:\n") - fmt.Printf(" cd %s\n", projectName) - fmt.Printf(" golanggraph multi-agent validate configs/%s\n", configFile) - fmt.Printf(" golanggraph multi-agent serve configs/%s\n", configFile) +// marshalConfig encodes a configuration in the requested format. +func marshalConfig(config interface{}, format string) ([]byte, error) { + switch format { + case "yaml", "yml": + return yaml.Marshal(config) + case "json": + return json.MarshalIndent(config, "", " ") + default: + return nil, fmt.Errorf("unsupported format %q", format) + } } // runMultiAgentValidate validates multi-agent configuration -func runMultiAgentValidate(args []string, strict, checkSchemas bool) { +func runMultiAgentValidate(out io.Writer, args []string, strict, checkSchemas bool) error { configFile := "configs/multi-agent.yaml" if len(args) > 0 { configFile = args[0] } - fmt.Printf("Validating multi-agent configuration: %s\n", configFile) - fmt.Printf("Strict mode: %t, Check schemas: %t\n", strict, checkSchemas) + _, _ = fmt.Fprintf(out, "Validating multi-agent configuration: %s\n", configFile) + _, _ = fmt.Fprintf(out, "Strict mode: %t, Check schemas: %t\n", strict, checkSchemas) - // Load configuration config, err := agent.LoadMultiAgentConfigFromFile(configFile) if err != nil { - fmt.Printf("Error loading configuration: %v\n", err) - os.Exit(1) + return err } - // Validate configuration if err := config.Validate(); err != nil { - fmt.Printf("Configuration validation failed: %v\n", err) - os.Exit(1) + return fmt.Errorf("configuration validation failed: %w", err) } - // Additional validations for strict mode + report := &validationReport{} + if checkSchemas { + schemaReport := validateAgentSchemas(config) + report.Errors = append(report.Errors, schemaReport.Errors...) + report.Warnings = append(report.Warnings, schemaReport.Warnings...) + } if strict { - if err := validateStrictMode(config); err != nil { - fmt.Printf("Strict validation failed: %v\n", err) - os.Exit(1) - } + report.Errors = append(report.Errors, validateStrictMode(config)...) } - // Schema validation - if checkSchemas { - if err := validateAgentSchemas(config); err != nil { - fmt.Printf("Schema validation failed: %v\n", err) - os.Exit(1) - } + for _, warning := range report.Warnings { + _, _ = fmt.Fprintf(out, " ⚠ %s\n", warning) + } + for _, problem := range report.Errors { + _, _ = fmt.Fprintf(out, " βœ— %s\n", problem) + } + if len(report.Errors) > 0 { + return fmt.Errorf("%s is invalid: %d problem(s)", configFile, len(report.Errors)) + } + if strict && len(report.Warnings) > 0 { + return fmt.Errorf("%s has %d warning(s) and --strict is set", configFile, len(report.Warnings)) } - fmt.Printf("βœ… Configuration validation passed!\n") - fmt.Printf("- Agents: %d\n", len(config.Agents)) - fmt.Printf("- Routing rules: %d\n", len(config.Routing.Rules)) - fmt.Printf("- Deployment type: %s\n", config.Deployment.Type) + _, _ = fmt.Fprintf(out, "βœ… Configuration validation passed!\n") + _, _ = fmt.Fprintf(out, "- Agents: %d\n", len(config.Agents)) + // Routing and Deployment are optional pointers, and printing through them + // unconditionally panicked on any configuration that omitted them -- after + // the success message had already been printed. + if config.Routing != nil { + _, _ = fmt.Fprintf(out, "- Routing rules: %d\n", len(config.Routing.Rules)) + } else { + _, _ = fmt.Fprintf(out, "- Routing rules: none configured\n") + } + if config.Deployment != nil && config.Deployment.Type != "" { + _, _ = fmt.Fprintf(out, "- Deployment type: %s\n", config.Deployment.Type) + } else { + _, _ = fmt.Fprintf(out, "- Deployment type: none configured\n") + } + return nil } // runMultiAgentDeploy deploys multiple agents -func runMultiAgentDeploy(args []string, deploymentType, environment string, dryRun, parallel bool) { +func runMultiAgentDeploy(out io.Writer, args []string, deploymentType, environment string, dryRun bool) error { configFile := "configs/multi-agent.yaml" if len(args) > 0 { configFile = args[0] } - fmt.Printf("Deploying multi-agent system: %s\n", configFile) - fmt.Printf("Type: %s, Environment: %s, Dry-run: %t, Parallel: %t\n", deploymentType, environment, dryRun, parallel) + _, _ = fmt.Fprintf(out, "Deploying multi-agent system: %s\n", configFile) + _, _ = fmt.Fprintf(out, "Type: %s, Environment: %s, Dry-run: %t\n", deploymentType, environment, dryRun) - // Load configuration config, err := agent.LoadMultiAgentConfigFromFile(configFile) if err != nil { - fmt.Printf("Error loading configuration: %v\n", err) - os.Exit(1) + return err } - - // Validate before deployment if err := config.Validate(); err != nil { - fmt.Printf("Configuration validation failed: %v\n", err) - os.Exit(1) + return fmt.Errorf("configuration validation failed: %w", err) } - // Override deployment type if specified + // config.Deployment is optional; assigning through it panicked. + if config.Deployment == nil { + config.Deployment = &agent.DeploymentConfig{} + } if deploymentType != "" { config.Deployment.Type = deploymentType } @@ -415,167 +484,157 @@ func runMultiAgentDeploy(args []string, deploymentType, environment string, dryR } if dryRun { - fmt.Printf("DRY RUN - Would deploy the following agents:\n") - for agentID, agentConfig := range config.Agents { - fmt.Printf(" - %s: %s (%s on %s)\n", agentID, agentConfig.Name, agentConfig.Type, agentConfig.Provider) + _, _ = fmt.Fprintf(out, "DRY RUN - would deploy the following agents:\n") + for _, agentID := range sortedAgentIDs(config) { + agentConfig := config.Agents[agentID] + _, _ = fmt.Fprintf(out, " - %s: %s (%s on %s)\n", agentID, agentConfig.Name, agentConfig.Type, agentConfig.Provider) } - return + return nil } - // Perform actual deployment switch config.Deployment.Type { - case "docker": - deployDocker(config, parallel) - case "kubernetes": - deployKubernetes(config, parallel) - case "serverless": - deployServerless(config, parallel) + case "docker", "kubernetes", "serverless": + // deployDocker/deployKubernetes/deployServerless printed "Deploying to + // Docker..." and returned without deploying anything, and the command + // exited 0. Generate the artifacts and say plainly that applying them + // is the operator's step. + return fmt.Errorf("deploying to %s is %w: generate the artifacts with 'golanggraph multi-agent generate %s' and apply them with your own tooling", + config.Deployment.Type, errNotImplemented, generateSubcommandFor(config.Deployment.Type)) default: - fmt.Printf("Unsupported deployment type: %s\n", config.Deployment.Type) - os.Exit(1) + return fmt.Errorf("unsupported deployment type %q (want docker, kubernetes or serverless)", config.Deployment.Type) } } +// generateSubcommandFor names the generator that produces artifacts for a +// deployment target. +func generateSubcommandFor(deploymentType string) string { + if deploymentType == "kubernetes" { + return "k8s" + } + return "docker" +} + +// sortedAgentIDs returns the agent IDs in a stable order; Go map iteration is +// randomized, so output ordering was different on every run. +func sortedAgentIDs(config *agent.MultiAgentConfig) []string { + ids := make([]string, 0, len(config.Agents)) + for id := range config.Agents { + ids = append(ids, id) + } + sort.Strings(ids) + return ids +} + // runMultiAgentServe starts multi-agent server -func runMultiAgentServe(args []string, host string, port int) { +func runMultiAgentServe(ctx context.Context, out io.Writer, args []string, host string, port int) error { configFile := "configs/multi-agent.yaml" if len(args) > 0 { configFile = args[0] } - fmt.Printf("Starting multi-agent server: %s\n", configFile) - fmt.Printf("Host: %s, Port: %d\n", host, port) - - // Load configuration - config, err := agent.LoadMultiAgentConfigFromFile(configFile) - if err != nil { - fmt.Printf("Error loading configuration: %v\n", err) - os.Exit(1) + if port <= 0 || port > 65535 { + return fmt.Errorf("invalid port %d", port) } - // Initialize components - llmManager := llm.NewProviderManager() - toolRegistry := tools.NewToolRegistry() - - // Setup LLM providers from shared config - if config.Shared != nil && config.Shared.LLMProviders != nil { - for providerName, providerConfig := range config.Shared.LLMProviders { - // Initialize provider based on type - // This is a simplified version - in a real implementation, - // you'd create the appropriate provider based on type - fmt.Printf("Setting up LLM provider: %s (%s)\n", providerName, providerConfig.Type) - } - } - - // Register default tools - toolRegistry.RegisterTool(tools.NewWebSearchTool()) - toolRegistry.RegisterTool(tools.NewCalculatorTool()) - toolRegistry.RegisterTool(tools.NewFileReadTool()) - toolRegistry.RegisterTool(tools.NewFileWriteTool()) - toolRegistry.RegisterTool(tools.NewShellTool()) - toolRegistry.RegisterTool(tools.NewHTTPTool()) - toolRegistry.RegisterTool(tools.NewTimeTool()) + _, _ = fmt.Fprintf(out, "Starting multi-agent server: %s\n", configFile) - // Create multi-agent manager - multiAgentManager, err := agent.NewMultiAgentManager(config, llmManager, toolRegistry) + config, err := agent.LoadMultiAgentConfigFromFile(configFile) if err != nil { - fmt.Printf("Error creating multi-agent manager: %v\n", err) - os.Exit(1) + return err } - // Start multi-agent manager - ctx := context.Background() - if err := multiAgentManager.Start(ctx); err != nil { - fmt.Printf("Error starting multi-agent manager: %v\n", err) - os.Exit(1) + // The previous implementation built a MultiAgentManager, started it, and + // then served an unrelated, empty server.NewServer -- while printing + // "Agent endpoints: http://host:port/agents". Nothing connected the two, so + // none of the advertised agent endpoints existed. Serve the agents through + // the auto-server, which generates an endpoint per registered agent. + autoServer := server.NewAutoServer(&server.AutoServerConfig{ + Host: host, + Port: port, + BasePath: "/api", + EnableWebUI: true, + EnablePlayground: true, + EnableSchemaAPI: true, + EnableMetricsAPI: true, + EnableCORS: true, + SchemaValidation: true, + ServerTimeout: 30 * time.Second, + MaxRequestSize: 10 * 1024 * 1024, + Middleware: []string{"cors", "logging", "recovery"}, + }) + + if err := autoServer.LoadAgentsFromConfig(configFile); err != nil { + return fmt.Errorf("failed to register agents: %w", err) } + _, _ = fmt.Fprintf(out, "Registered %d agent(s): %s\n", len(config.Agents), strings.Join(sortedAgentIDs(config), ", ")) - // Create server configuration - serverConfig := &server.ServerConfig{ - Host: host, - Port: port, - ReadTimeout: 30 * time.Second, - WriteTimeout: 30 * time.Second, - MaxHeaderBytes: 1 << 20, - EnableCORS: true, - StaticDir: "./static", - DevMode: viper.GetBool("dev-mode"), + if err := checkAddressAvailable(host, port); err != nil { + return err } - // Create and start server with multi-agent router - srv := server.NewServer(serverConfig) - // Note: The server will use its own router, multi-agent manager handles routing internally - - fmt.Printf("Multi-agent server started on %s:%d\n", host, port) - fmt.Printf("Health check: http://%s:%d/health\n", host, port) - fmt.Printf("Agent endpoints: http://%s:%d/agents\n", host, port) - fmt.Printf("Metrics: http://%s:%d/metrics\n", host, port) + _, _ = fmt.Fprintf(out, "Multi-agent server listening on %s:%d\n", host, port) + _, _ = fmt.Fprintf(out, "Health check: http://%s:%d/health\n", host, port) + _, _ = fmt.Fprintf(out, "Agent endpoints: http://%s:%d/agents\n", host, port) - if err := srv.Start(); err != nil { - fmt.Printf("Server failed to start: %v\n", err) - os.Exit(1) + if err := autoServer.Start(ctx); err != nil { + return fmt.Errorf("server failed: %w", err) } + return nil } -// runMultiAgentStatus checks status of deployed agents -func runMultiAgentStatus(args []string, outputFormat string, watch bool) { +// runMultiAgentStatus reports the configured agents. +func runMultiAgentStatus(out io.Writer, args []string, outputFormat string, watch bool) error { configFile := "configs/multi-agent.yaml" if len(args) > 0 { configFile = args[0] } - fmt.Printf("Checking multi-agent status: %s\n", configFile) + if watch { + // --watch used to fall into `select {}`, blocking forever after saying + // it was watching for changes. Nothing was ever polled or printed. + return fmt.Errorf("--watch is %w for multi-agent status", errNotImplemented) + } - // Load configuration config, err := agent.LoadMultiAgentConfigFromFile(configFile) if err != nil { - fmt.Printf("Error loading configuration: %v\n", err) - os.Exit(1) + return err } - // In a real implementation, this would connect to the running multi-agent system - // and fetch actual status information - - status := map[string]interface{}{ - "timestamp": time.Now(), - "config": configFile, - "agents": make(map[string]interface{}), - } + _, _ = fmt.Fprintf(out, "Configured agents in %s (this reads the configuration; it does not contact a running deployment):\n", configFile) + agents := make(map[string]interface{}, len(config.Agents)) for agentID, agentConfig := range config.Agents { - status["agents"].(map[string]interface{})[agentID] = map[string]interface{}{ - "name": agentConfig.Name, - "type": agentConfig.Type, - "status": "unknown", // Would be fetched from actual deployment - "health_status": "unknown", - "request_count": 0, - "error_count": 0, + agents[agentID] = map[string]interface{}{ + "name": agentConfig.Name, + "type": agentConfig.Type, + "provider": agentConfig.Provider, + "model": agentConfig.Model, } } + status := map[string]interface{}{ + "timestamp": time.Now().UTC().Format(time.RFC3339), + "config": configFile, + "agents": agents, + } - // Output status switch outputFormat { - case "json": - output, _ := json.MarshalIndent(status, "", " ") - fmt.Println(string(output)) - case "yaml": - output, _ := yaml.Marshal(status) - fmt.Println(string(output)) - default: - // Table format - fmt.Printf("\nAgent Status:\n") - fmt.Printf("%-20s %-15s %-10s %-10s\n", "Agent ID", "Type", "Status", "Health") - fmt.Printf("%-20s %-15s %-10s %-10s\n", "--------", "----", "------", "------") - for agentID, agentConfig := range config.Agents { - fmt.Printf("%-20s %-15s %-10s %-10s\n", agentID, agentConfig.Type, "unknown", "unknown") + case "json", "yaml": + encoded, err := marshalConfig(status, outputFormat) + if err != nil { + return err } + _, _ = fmt.Fprintln(out, string(encoded)) + case "table", "": + _, _ = fmt.Fprintf(out, "\n%-20s %-15s %-12s %-20s\n", "Agent ID", "Type", "Provider", "Model") + _, _ = fmt.Fprintf(out, "%-20s %-15s %-12s %-20s\n", "--------", "----", "--------", "-----") + for _, agentID := range sortedAgentIDs(config) { + agentConfig := config.Agents[agentID] + _, _ = fmt.Fprintf(out, "%-20s %-15s %-12s %-20s\n", agentID, agentConfig.Type, agentConfig.Provider, agentConfig.Model) + } + default: + return fmt.Errorf("unsupported output format %q (want table, json or yaml)", outputFormat) } - - if watch { - fmt.Printf("\nWatching for status changes (press Ctrl+C to exit)...\n") - // In a real implementation, this would watch for status changes - select {} // Block forever - } + return nil } // Helper functions for generating project artifacts @@ -750,29 +809,21 @@ func setupRoutingRules(config *agent.MultiAgentConfig, routingType string) { } } -func createIndividualAgentConfigs(projectName string, config *agent.MultiAgentConfig, format string) { +func createIndividualAgentConfigs(projectName string, config *agent.MultiAgentConfig, format string) error { for agentID, agentConfig := range config.Agents { - agentDir := filepath.Join(projectName, "agents", agentID) - configFile := fmt.Sprintf("config.%s", format) - configPath := filepath.Join(agentDir, configFile) - - var configData []byte - var err error - - switch format { - case "yaml", "yml": - configData, err = yaml.Marshal(agentConfig) - case "json": - configData, err = json.MarshalIndent(agentConfig, "", " ") + configData, err := marshalConfig(agentConfig, format) + if err != nil { + return fmt.Errorf("agent %s: %w", agentID, err) } - - if err == nil { - os.WriteFile(configPath, configData, 0600) + path := filepath.Join(projectName, "agents", agentID, fmt.Sprintf("config.%s", format)) + if err := writeFileChecked(path, string(configData)); err != nil { + return err } } + return nil } -func createDockerComposeFile(projectName string, config *agent.MultiAgentConfig) { +func createDockerComposeFile(projectName string, config *agent.MultiAgentConfig) error { dockerCompose := `version: '3.8' services: multi-agent: @@ -816,11 +867,10 @@ volumes: ollama_data: ` - composePath := filepath.Join(projectName, "docker-compose.yml") - os.WriteFile(composePath, []byte(dockerCompose), 0600) + return writeFileChecked(filepath.Join(projectName, "docker-compose.yml"), dockerCompose) } -func createK8sManifests(projectName string, config *agent.MultiAgentConfig) { +func createK8sManifests(projectName string, config *agent.MultiAgentConfig) error { k8sDir := filepath.Join(projectName, "k8s") // Deployment manifest @@ -868,11 +918,13 @@ spec: type: LoadBalancer ` - os.WriteFile(filepath.Join(k8sDir, "deployment.yaml"), []byte(deployment), 0600) - os.WriteFile(filepath.Join(k8sDir, "service.yaml"), []byte(service), 0600) + if err := writeFileChecked(filepath.Join(k8sDir, "deployment.yaml"), deployment); err != nil { + return err + } + return writeFileChecked(filepath.Join(k8sDir, "service.yaml"), service) } -func createProjectREADME(projectName string, config *agent.MultiAgentConfig) { +func createProjectREADME(projectName string, config *agent.MultiAgentConfig) error { readme := fmt.Sprintf("# %s\n\nMulti-agent GoLangGraph project with %d agents.\n\n", projectName, len(config.Agents)) readme += "## Quick Start\n\n" @@ -897,7 +949,8 @@ func createProjectREADME(projectName string, config *agent.MultiAgentConfig) { readme += "The multi-agent configuration is defined in `configs/multi-agent.yaml`.\n\n" readme += "### Agents\n\n" - for agentID, agentConfig := range config.Agents { + for _, agentID := range sortedAgentIDs(config) { + agentConfig := config.Agents[agentID] readme += fmt.Sprintf("- **%s**: %s (%s)\n", agentID, agentConfig.Name, agentConfig.Type) } @@ -937,46 +990,304 @@ func createProjectREADME(projectName string, config *agent.MultiAgentConfig) { readme += "- Metrics: `http://localhost:8080/metrics`\n" readme += "- Agent Status: `http://localhost:8080/agents`\n" - readmePath := filepath.Join(projectName, "README.md") - os.WriteFile(readmePath, []byte(readme), 0600) + return writeFileChecked(filepath.Join(projectName, "README.md"), readme) } -// Additional validation functions -func validateStrictMode(config *agent.MultiAgentConfig) error { - // Add strict validation logic here - return nil +// validateStrictMode applies the checks that MultiAgentConfig.Validate does not. +// +// It used to be `return nil` with a "// Add strict validation logic here" +// comment, so "multi-agent validate --strict" reported that strict validation +// had passed without performing any. +func validateStrictMode(config *agent.MultiAgentConfig) []string { + var problems []string + + for _, agentID := range sortedAgentIDs(config) { + agentConfig := config.Agents[agentID] + if agentConfig.SystemPrompt == "" { + problems = append(problems, fmt.Sprintf("agent %s: no system prompt", agentID)) + } + if agentConfig.ID != "" && agentConfig.ID != agentID { + problems = append(problems, fmt.Sprintf("agent %s: id field is %q, which does not match its key", agentID, agentConfig.ID)) + } + if !knownProviders[strings.ToLower(agentConfig.Provider)] { + problems = append(problems, fmt.Sprintf("agent %s: provider %q is not one of the built-in providers", agentID, agentConfig.Provider)) + } + } + + if config.Routing != nil { + if config.Routing.DefaultAgent == "" { + problems = append(problems, "routing: no default agent") + } + patterns := map[string]string{} + for _, rule := range config.Routing.Rules { + if rule.Pattern == "" { + problems = append(problems, fmt.Sprintf("routing: rule %s has no pattern", rule.ID)) + continue + } + if previous, clash := patterns[rule.Pattern]; clash { + problems = append(problems, fmt.Sprintf("routing: pattern %q is claimed by both %s and %s", rule.Pattern, previous, rule.AgentID)) + } + patterns[rule.Pattern] = rule.AgentID + } + for _, agentID := range sortedAgentIDs(config) { + routed := false + for _, rule := range config.Routing.Rules { + if rule.AgentID == agentID { + routed = true + break + } + } + if !routed && config.Routing.DefaultAgent != agentID { + problems = append(problems, fmt.Sprintf("agent %s: no routing rule reaches it", agentID)) + } + } + } + + return problems } -func validateAgentSchemas(config *agent.MultiAgentConfig) error { - // Add schema validation logic here +// validateAgentSchemas checks each agent definition against what the runtime +// requires: a known type, resolvable tools and a graph that builds. +// +// This was `return nil` too, while --check-schemas defaults to true -- so every +// "multi-agent validate" run reported schema validation it never did. +func validateAgentSchemas(config *agent.MultiAgentConfig) *validationReport { + report := &validationReport{} + + toolRegistry := tools.NewToolRegistry() + known := map[string]bool{} + for _, name := range toolRegistry.ListTools() { + known[name] = true + } + llmManager := llm.NewProviderManager() + + for _, agentID := range sortedAgentIDs(config) { + // Validate a copy with the framework defaults filled in for keys the + // file omits, so a configuration that merely leaves max_tokens unset is + // not reported as broken (the runtime defaults it too). + agentConfig := *config.Agents[agentID] + defaults := agent.DefaultAgentConfig() + if agentConfig.MaxTokens == 0 { + agentConfig.MaxTokens = defaults.MaxTokens + } + if agentConfig.MaxIterations == 0 { + agentConfig.MaxIterations = defaults.MaxIterations + } + + if err := agentConfig.Validate(); err != nil { + report.errorf("agent %s: %v", agentID, err) + continue + } + + switch agentConfig.Type { + case agent.AgentTypeChat, agent.AgentTypeReAct, agent.AgentTypeTool: + default: + // The runtime silently falls back to a chat graph for an unknown + // type, so an operator would never learn of the typo. + report.errorf("agent %s: unknown type %q (want chat, react or tool)", agentID, agentConfig.Type) + continue + } + + for _, tool := range agentConfig.Tools { + if !known[tool] { + // A warning, not an error: tools can also be registered by the + // operator's own Go code at start-up. + report.warnf("agent %s: tool %q is not registered", agentID, tool) + } + } + + built := agent.NewAgent(&agentConfig, llmManager, toolRegistry) + graph := built.GetGraph() + if graph == nil { + report.errorf("agent %s: no execution graph was built", agentID) + continue + } + if err := graph.Validate(); err != nil { + report.errorf("agent %s: execution graph is invalid: %v", agentID, err) + } + } + + return report +} + +// runGenerateDocker writes the Docker deployment artifacts. +// +// It used to print "Generating Docker deployment files..." and generate +// nothing, ignoring --output and --multi-service entirely. +func runGenerateDocker(out io.Writer, args []string, outputDir string, multiService bool) error { + configFile := "configs/multi-agent.yaml" + if len(args) > 0 { + configFile = args[0] + } + if outputDir == "" { + outputDir = "./deploy" + } + + config, err := agent.LoadMultiAgentConfigFromFile(configFile) + if err != nil { + return err + } + if err := os.MkdirAll(outputDir, 0750); err != nil { + return fmt.Errorf("failed to create %s: %w", outputDir, err) + } + + compose := dockerComposeFor(config, multiService) + composePath := filepath.Join(outputDir, "docker-compose.yml") + if err := writeFileChecked(composePath, compose); err != nil { + return err + } + + dockerfilePath := filepath.Join(outputDir, "Dockerfile") + if err := writeFileChecked(dockerfilePath, agentDockerfile); err != nil { + return err + } + + _, _ = fmt.Fprintf(out, "Generated %s\n", composePath) + _, _ = fmt.Fprintf(out, "Generated %s\n", dockerfilePath) return nil } -// Deployment functions -func deployDocker(config *agent.MultiAgentConfig, parallel bool) { - fmt.Printf("Deploying to Docker...\n") - // Implementation for Docker deployment +// dockerComposeFor renders a compose file for the configured agents. With +// --multi-service each agent gets its own service and port. +func dockerComposeFor(config *agent.MultiAgentConfig, multiService bool) string { + var b strings.Builder + b.WriteString("# Generated by golanggraph multi-agent generate docker\n") + b.WriteString("services:\n") + + if multiService { + port := 8080 + for _, agentID := range sortedAgentIDs(config) { + _, _ = fmt.Fprintf(&b, " %s:\n", agentID) + b.WriteString(" build: .\n") + _, _ = fmt.Fprintf(&b, " command: [\"serve\", \"--host\", \"0.0.0.0\", \"--port\", \"8080\"]\n") + _, _ = fmt.Fprintf(&b, " ports:\n - \"%d:8080\"\n", port) + _, _ = fmt.Fprintf(&b, " environment:\n - GOLANGGRAPH_AGENT_ID=%s\n", agentID) + b.WriteString(" restart: unless-stopped\n") + port++ + } + } else { + b.WriteString(" multi-agent:\n") + b.WriteString(" build: .\n") + b.WriteString(" ports:\n - \"8080:8080\"\n") + b.WriteString(" volumes:\n - ./configs:/app/configs:ro\n") + b.WriteString(" restart: unless-stopped\n") + } + + return b.String() } -func deployKubernetes(config *agent.MultiAgentConfig, parallel bool) { - fmt.Printf("Deploying to Kubernetes...\n") - // Implementation for Kubernetes deployment +// runGenerateK8s writes the Kubernetes manifests. +// +// It used to print "Generating Kubernetes manifests..." and generate nothing, +// ignoring --output and --namespace entirely. +func runGenerateK8s(out io.Writer, args []string, outputDir, namespace string) error { + configFile := "configs/multi-agent.yaml" + if len(args) > 0 { + configFile = args[0] + } + if outputDir == "" { + outputDir = "./k8s" + } + if namespace == "" { + namespace = "golanggraph" + } + + config, err := agent.LoadMultiAgentConfigFromFile(configFile) + if err != nil { + return err + } + if err := os.MkdirAll(outputDir, 0750); err != nil { + return fmt.Errorf("failed to create %s: %w", outputDir, err) + } + + replicas := 1 + if config.Deployment != nil && config.Deployment.Replicas > 0 { + replicas = config.Deployment.Replicas + } + + manifests := map[string]string{ + "namespace.yaml": k8sNamespaceManifest(namespace), + "deployment.yaml": k8sDeploymentManifest(namespace, replicas), + "service.yaml": k8sServiceManifest(namespace), + } + + names := make([]string, 0, len(manifests)) + for name := range manifests { + names = append(names, name) + } + sort.Strings(names) + + for _, name := range names { + path := filepath.Join(outputDir, name) + if err := writeFileChecked(path, manifests[name]); err != nil { + return err + } + _, _ = fmt.Fprintf(out, "Generated %s\n", path) + } + return nil } -func deployServerless(config *agent.MultiAgentConfig, parallel bool) { - fmt.Printf("Deploying to serverless platform...\n") - // Implementation for serverless deployment +func k8sNamespaceManifest(namespace string) string { + return fmt.Sprintf(`apiVersion: v1 +kind: Namespace +metadata: + name: %s +`, namespace) } -// Generate functions -func runGenerateDocker(args []string, outputDir string, multiService bool) { - fmt.Printf("Generating Docker deployment files...\n") - // Implementation for generating Docker files +func k8sDeploymentManifest(namespace string, replicas int) string { + return fmt.Sprintf(`apiVersion: apps/v1 +kind: Deployment +metadata: + name: golanggraph-multi-agent + namespace: %s + labels: + app: golanggraph-multi-agent +spec: + replicas: %d + selector: + matchLabels: + app: golanggraph-multi-agent + template: + metadata: + labels: + app: golanggraph-multi-agent + spec: + containers: + - name: multi-agent + image: golanggraph-multi-agent:latest + ports: + - containerPort: 8080 + readinessProbe: + httpGet: + path: /health + port: 8080 + initialDelaySeconds: 5 + periodSeconds: 5 + livenessProbe: + httpGet: + path: /health + port: 8080 + initialDelaySeconds: 30 + periodSeconds: 30 +`, namespace, replicas) } -func runGenerateK8s(args []string, outputDir, namespace string) { - fmt.Printf("Generating Kubernetes manifests...\n") - // Implementation for generating K8s manifests +func k8sServiceManifest(namespace string) string { + return fmt.Sprintf(`apiVersion: v1 +kind: Service +metadata: + name: golanggraph-multi-agent + namespace: %s +spec: + selector: + app: golanggraph-multi-agent + ports: + - protocol: TCP + port: 80 + targetPort: 8080 + type: ClusterIP +`, namespace) } // runMultiAgentLoad loads agent definitions from Go files or plugins @@ -986,91 +1297,107 @@ func runMultiAgentLoad(cmd *cobra.Command, args []string) error { path = args[0] } - recursive, _ := cmd.Flags().GetBool("recursive") - include, _ := cmd.Flags().GetStringSlice("include") - exclude, _ := cmd.Flags().GetStringSlice("exclude") - validate, _ := cmd.Flags().GetBool("validate") - verbose, _ := cmd.Flags().GetBool("verbose") + validate, err := cmd.Flags().GetBool("validate") + if err != nil { + return err + } + verbose, err := cmd.Flags().GetBool("verbose") + if err != nil { + return err + } - fmt.Printf("Loading agent definitions from: %s\n", path) + out := cmd.OutOrStdout() + _, _ = fmt.Fprintf(out, "Loading agent definitions from: %s\n", path) registry := agent.GetGlobalRegistry() - // Check if it's a plugin file - if strings.HasSuffix(path, ".so") { - if verbose { - fmt.Printf("Loading plugin: %s\n", path) - } - - if err := registry.LoadFromPlugin(path); err != nil { - return fmt.Errorf("failed to load plugin: %w", err) - } - - fmt.Printf("Successfully loaded plugin: %s\n", path) - } else { - // Load from directory - if verbose { - fmt.Printf("Scanning directory for Go files...\n") - fmt.Printf("Include patterns: %v\n", include) - fmt.Printf("Exclude patterns: %v\n", exclude) - fmt.Printf("Recursive: %v\n", recursive) - } + // Only plugin loading exists. Directory loading used to print "not yet + // implemented" and then return nil, so a script that checked the exit + // status was told the load had succeeded. + if !strings.HasSuffix(path, ".so") { + return fmt.Errorf("loading agent definitions from a directory is %w; build your agents as a Go plugin and pass the .so file", errNotImplemented) + } - // In a real implementation, this would scan the directory - // for Go files and load agent definitions - fmt.Printf("Directory-based loading not yet implemented\n") - fmt.Printf("Please use plugin-based loading instead\n") - return nil + if verbose { + _, _ = fmt.Fprintf(out, "Loading plugin: %s\n", path) } + if err := registry.LoadFromPlugin(path); err != nil { + return fmt.Errorf("failed to load plugin: %w", err) + } + _, _ = fmt.Fprintf(out, "Successfully loaded plugin: %s\n", path) // List loaded agents definitions := registry.ListDefinitions() factories := registry.ListFactories() + sort.Strings(definitions) + sort.Strings(factories) - fmt.Printf("\nLoaded agents:\n") - fmt.Printf(" Definitions: %d\n", len(definitions)) - fmt.Printf(" Factories: %d\n", len(factories)) + _, _ = fmt.Fprintf(out, "\nLoaded agents:\n") + _, _ = fmt.Fprintf(out, " Definitions: %d\n", len(definitions)) + _, _ = fmt.Fprintf(out, " Factories: %d\n", len(factories)) - if verbose { - fmt.Printf("\nDefinitions: %v\n", definitions) - fmt.Printf("Factories: %v\n", factories) + if !validate { + return nil } - // Validate if requested - if validate { - fmt.Printf("\nValidating loaded agent definitions...\n") + _, _ = fmt.Fprintf(out, "\nValidating loaded agent definitions...\n") - for _, defID := range definitions { - if def, exists := registry.GetDefinition(defID); exists { - if err := def.Validate(); err != nil { - fmt.Printf(" ❌ %s: %v\n", defID, err) - } else { - fmt.Printf(" βœ… %s: valid\n", defID) - } - } + var invalid []string + for _, defID := range definitions { + def, exists := registry.GetDefinition(defID) + if !exists { + continue + } + if err := def.Validate(); err != nil { + _, _ = fmt.Fprintf(out, " ❌ %s: %v\n", defID, err) + invalid = append(invalid, defID) + continue } + _, _ = fmt.Fprintf(out, " βœ… %s: valid\n", defID) + } - for _, factoryID := range factories { - // Create temporary instance to validate - factory := registry.ListFactories() - if len(factory) > 0 { - fmt.Printf(" βœ… Factory %s: valid\n", factoryID) - } + for _, factoryID := range factories { + // This used to print "βœ… Factory X: valid" whenever the registry held + // any factory at all, without ever looking at the factory in question. + // Build the definition it produces and validate that. + definition, err := registry.CreateAgentFromFactory(factoryID, llm.NewProviderManager(), tools.NewToolRegistry()) + if err != nil || definition == nil { + _, _ = fmt.Fprintf(out, " ❌ Factory %s: %v\n", factoryID, err) + invalid = append(invalid, factoryID) + continue } + _, _ = fmt.Fprintf(out, " βœ… Factory %s: valid\n", factoryID) } + if len(invalid) > 0 { + return fmt.Errorf("%d agent definition(s) are invalid: %s", len(invalid), strings.Join(invalid, ", ")) + } return nil } // runMultiAgentList lists all registered agent definitions func runMultiAgentList(cmd *cobra.Command, args []string) error { - format, _ := cmd.Flags().GetString("format") - filter, _ := cmd.Flags().GetString("filter") - showMetadata, _ := cmd.Flags().GetBool("show-metadata") - showConfig, _ := cmd.Flags().GetBool("show-config") + format, err := cmd.Flags().GetString("format") + if err != nil { + return err + } + filter, err := cmd.Flags().GetString("filter") + if err != nil { + return err + } + showMetadata, err := cmd.Flags().GetBool("show-metadata") + if err != nil { + return err + } + showConfig, err := cmd.Flags().GetBool("show-config") + if err != nil { + return err + } + out := cmd.OutOrStdout() registry := agent.GetGlobalRegistry() infos := registry.GetAgentInfo() + sort.Slice(infos, func(i, j int) bool { return infos[i].ID < infos[j].ID }) // Apply filter if specified if filter != "" { @@ -1089,20 +1416,19 @@ func runMultiAgentList(cmd *cobra.Command, args []string) error { if err != nil { return fmt.Errorf("failed to marshal JSON: %w", err) } - fmt.Println(string(output)) + _, _ = fmt.Fprintln(out, string(output)) case "yaml": output, err := yaml.Marshal(infos) if err != nil { return fmt.Errorf("failed to marshal YAML: %w", err) } - fmt.Println(string(output)) + _, _ = fmt.Fprintln(out, string(output)) - default: - // Table format - fmt.Printf("Agent Definitions (%d total):\n\n", len(infos)) - fmt.Printf("%-20s %-12s %-15s %-10s\n", "ID", "Source", "Type", "Model") - fmt.Printf("%-20s %-12s %-15s %-10s\n", "--", "------", "----", "-----") + case "table", "": + _, _ = fmt.Fprintf(out, "Agent Definitions (%d total):\n\n", len(infos)) + _, _ = fmt.Fprintf(out, "%-20s %-12s %-15s %-10s\n", "ID", "Source", "Type", "Model") + _, _ = fmt.Fprintf(out, "%-20s %-12s %-15s %-10s\n", "--", "------", "----", "-----") for _, info := range infos { model := "N/A" @@ -1113,22 +1439,30 @@ func runMultiAgentList(cmd *cobra.Command, args []string) error { agentType = string(info.Config.Type) } - fmt.Printf("%-20s %-12s %-15s %-10s\n", + _, _ = fmt.Fprintf(out, "%-20s %-12s %-15s %-10s\n", info.ID, info.Source, agentType, model) if showConfig && info.Config != nil { - fmt.Printf(" Config: Name=%s, Provider=%s, Tools=%v\n", + _, _ = fmt.Fprintf(out, " Config: Name=%s, Provider=%s, Tools=%v\n", info.Config.Name, info.Config.Provider, info.Config.Tools) } if showMetadata && len(info.Metadata) > 0 { - fmt.Printf(" Metadata: ") - for k, v := range info.Metadata { - fmt.Printf("%s=%v ", k, v) + keys := make([]string, 0, len(info.Metadata)) + for k := range info.Metadata { + keys = append(keys, k) + } + sort.Strings(keys) + _, _ = fmt.Fprintf(out, " Metadata: ") + for _, k := range keys { + _, _ = fmt.Fprintf(out, "%s=%v ", k, info.Metadata[k]) } - fmt.Printf("\n") + _, _ = fmt.Fprintf(out, "\n") } } + + default: + return fmt.Errorf("unsupported output format %q (want table, json or yaml)", format) } return nil diff --git a/cmd/golanggraph/multi_agent_commands_test.go b/cmd/golanggraph/multi_agent_commands_test.go new file mode 100644 index 0000000..2d12b47 --- /dev/null +++ b/cmd/golanggraph/multi_agent_commands_test.go @@ -0,0 +1,518 @@ +// Copyright (c) 2024 GoLangGraph Team +// +// Licensed under the MIT License. See LICENSE file in the project root for full license information. + +package main + +import ( + "bytes" + "context" + "encoding/json" + "os" + "path/filepath" + "strings" + "testing" + "time" + + "github.com/spf13/cobra" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// A multi-agent file with no routing and no deployment section: both are +// optional pointers in MultiAgentConfig. +const minimalMultiAgentYAML = `name: minimal +agents: + a1: + id: a1 + name: A1 + type: chat + model: gpt-3.5-turbo + provider: openai + systemprompt: "You are A1." + maxtokens: 1000 +` + +// A complete file: note that agent fields are spelled the way gopkg.in/yaml.v3 +// maps them onto agent.AgentConfig, which carries JSON tags only. +const fullMultiAgentYAML = `name: fleet +version: "1.0.0" +agents: + alpha: + id: alpha + name: Alpha + type: chat + model: gpt-4 + provider: openai + systemprompt: "You are Alpha." + maxtokens: 1000 + tools: + - calculator + beta: + id: beta + name: Beta + type: react + model: gpt-4 + provider: openai + systemprompt: "You are Beta." + maxtokens: 1000 +routing: + type: path + default_agent: alpha + rules: + - id: rule-1 + pattern: /alpha + agent_id: alpha + method: POST + - id: rule-2 + pattern: /beta + agent_id: beta + method: POST +deployment: + type: docker + environment: development + replicas: 3 +` + +// Regression: runMultiAgentValidate printed "βœ… Configuration validation +// passed!" and then dereferenced config.Routing and config.Deployment +// unconditionally, so any configuration without those optional sections +// panicked with a nil pointer dereference after reporting success. +func TestMultiAgentValidate_ConfigWithoutRoutingOrDeployment(t *testing.T) { + path := writeTestFile(t, t.TempDir(), "minimal.yaml", minimalMultiAgentYAML) + + var out bytes.Buffer + require.NotPanics(t, func() { + require.NoError(t, runMultiAgentValidate(&out, []string{path}, false, true)) + }) + assert.Contains(t, out.String(), "Routing rules: none configured") + assert.Contains(t, out.String(), "Deployment type: none configured") +} + +func TestMultiAgentValidate_AcceptsACompleteConfiguration(t *testing.T) { + path := writeTestFile(t, t.TempDir(), "multi.yaml", fullMultiAgentYAML) + + var out bytes.Buffer + require.NoError(t, runMultiAgentValidate(&out, []string{path}, true, true)) + assert.Contains(t, out.String(), "Agents: 2") + assert.Contains(t, out.String(), "Routing rules: 2") +} + +func TestMultiAgentValidate_MissingAndMalformedFiles(t *testing.T) { + dir := t.TempDir() + + require.Error(t, runMultiAgentValidate(&bytes.Buffer{}, []string{filepath.Join(dir, "absent.yaml")}, false, true)) + + broken := writeTestFile(t, dir, "broken.yaml", "agents: [\n") + require.Error(t, runMultiAgentValidate(&bytes.Buffer{}, []string{broken}, false, true)) + + empty := writeTestFile(t, dir, "noagents.yaml", "name: nobody\n") + err := runMultiAgentValidate(&bytes.Buffer{}, []string{empty}, false, true) + require.Error(t, err) + assert.Contains(t, err.Error(), "at least one agent") +} + +// Regression: validateStrictMode was `return nil` with a "add strict validation +// logic here" comment, so --strict reported that strict validation had passed +// without performing any. +func TestMultiAgentValidate_StrictFindsRealProblems(t *testing.T) { + path := writeTestFile(t, t.TempDir(), "multi.yaml", `name: fleet +agents: + alpha: + id: alpha + name: Alpha + type: chat + model: gpt-4 + provider: openai + maxtokens: 1000 + beta: + id: wrong-id + name: Beta + type: chat + model: gpt-4 + provider: mystery-inc + systemprompt: "hi" + maxtokens: 1000 +routing: + type: path + rules: + - id: rule-1 + pattern: /same + agent_id: alpha + - id: rule-2 + pattern: /same + agent_id: beta +`) + + var out bytes.Buffer + require.NoError(t, runMultiAgentValidate(&out, []string{path}, false, true), + "these are strict-mode problems only") + + out.Reset() + err := runMultiAgentValidate(&out, []string{path}, true, true) + require.Error(t, err) + + reported := out.String() + assert.Contains(t, reported, "no system prompt") + assert.Contains(t, reported, "does not match its key") + assert.Contains(t, reported, "not one of the built-in providers") + assert.Contains(t, reported, "no default agent") + assert.Contains(t, reported, `pattern "/same" is claimed by both`) +} + +// Regression: validateAgentSchemas was `return nil` while --check-schemas +// defaults to true, so every run reported schema validation it never did. +func TestMultiAgentValidate_SchemaCheckFindsUnknownToolsAndTypes(t *testing.T) { + dir := t.TempDir() + badType := writeTestFile(t, dir, "badtype.yaml", `name: fleet +agents: + beta: + id: beta + name: Beta + type: reactt + model: gpt-4 + provider: openai + systemprompt: "hi" + maxtokens: 1000 +`) + unknownTool := writeTestFile(t, dir, "tool.yaml", `name: fleet +agents: + alpha: + id: alpha + name: Alpha + type: chat + model: gpt-4 + provider: openai + systemprompt: "hi" + maxtokens: 1000 + tools: + - no_such_tool +`) + + var out bytes.Buffer + require.NoError(t, runMultiAgentValidate(&out, []string{badType}, false, false), + "--check-schemas=false must not run the schema checks") + + out.Reset() + err := runMultiAgentValidate(&out, []string{badType}, false, true) + require.Error(t, err, "an agent type the runtime does not implement must be reported") + assert.Contains(t, out.String(), `unknown type "reactt"`) + + // A tool the CLI does not know may still be registered by the operator's + // own code, so it is a warning that --strict turns into a failure. + out.Reset() + require.NoError(t, runMultiAgentValidate(&out, []string{unknownTool}, false, true)) + assert.Contains(t, out.String(), `tool "no_such_tool" is not registered`) + require.Error(t, runMultiAgentValidate(&bytes.Buffer{}, []string{unknownTool}, true, true)) +} + +// A configuration that merely leaves optional keys unset is not broken: the +// runtime defaults them, and so must the validator. +func TestMultiAgentValidate_OmittedOptionalKeysAreDefaulted(t *testing.T) { + path := writeTestFile(t, t.TempDir(), "multi.yaml", `name: fleet +agents: + alpha: + id: alpha + name: Alpha + type: chat + model: gpt-4 + provider: openai + systemprompt: "hi" +`) + + require.NoError(t, runMultiAgentValidate(&bytes.Buffer{}, []string{path}, false, true)) +} + +// Regression: deployDocker/deployKubernetes/deployServerless printed +// "Deploying to Docker..." and returned without deploying anything, and the +// command exited 0. +func TestMultiAgentDeploy_DoesNotClaimAnUndoneDeployment(t *testing.T) { + path := writeTestFile(t, t.TempDir(), "multi.yaml", fullMultiAgentYAML) + + var out bytes.Buffer + err := runMultiAgentDeploy(&out, []string{path}, "docker", "production", false) + + require.Error(t, err, "a deployment that did not happen must not exit zero") + assert.ErrorIs(t, err, errNotImplemented) + assert.NotContains(t, out.String(), "Deploying to Docker...") +} + +func TestMultiAgentDeploy_DryRunListsAgentsAndSucceeds(t *testing.T) { + path := writeTestFile(t, t.TempDir(), "multi.yaml", fullMultiAgentYAML) + + var out bytes.Buffer + require.NoError(t, runMultiAgentDeploy(&out, []string{path}, "docker", "development", true)) + + listing := out.String() + assert.Contains(t, listing, "alpha") + assert.Contains(t, listing, "beta") + assert.Less(t, strings.Index(listing, "alpha"), strings.Index(listing, "beta"), + "agents must be listed in a stable order") +} + +// Regression: `config.Deployment.Type = deploymentType` panicked whenever the +// configuration had no deployment section. +func TestMultiAgentDeploy_ConfigWithoutDeploymentSection(t *testing.T) { + path := writeTestFile(t, t.TempDir(), "minimal.yaml", minimalMultiAgentYAML) + + require.NotPanics(t, func() { + require.NoError(t, runMultiAgentDeploy(&bytes.Buffer{}, []string{path}, "docker", "", true)) + }) +} + +func TestMultiAgentDeploy_UnknownDeploymentTypeIsRejected(t *testing.T) { + path := writeTestFile(t, t.TempDir(), "multi.yaml", fullMultiAgentYAML) + + err := runMultiAgentDeploy(&bytes.Buffer{}, []string{path}, "carrier-pigeon", "", false) + require.Error(t, err) + assert.Contains(t, err.Error(), "unsupported deployment type") +} + +// Regression: runGenerateDocker printed "Generating Docker deployment files..." +// and generated nothing, ignoring --output and --multi-service entirely. +func TestMultiAgentGenerateDocker_WritesTheFiles(t *testing.T) { + dir := t.TempDir() + path := writeTestFile(t, dir, "multi.yaml", fullMultiAgentYAML) + outputDir := filepath.Join(dir, "deploy") + + var out bytes.Buffer + require.NoError(t, runGenerateDocker(&out, []string{path}, outputDir, true)) + + compose, err := os.ReadFile(filepath.Join(outputDir, "docker-compose.yml")) + require.NoError(t, err, "the compose file must actually exist") + assert.Contains(t, string(compose), "alpha:") + assert.Contains(t, string(compose), "beta:") + + dockerfile, err := os.ReadFile(filepath.Join(outputDir, "Dockerfile")) + require.NoError(t, err) + assert.Contains(t, string(dockerfile), "FROM golang") +} + +func TestMultiAgentGenerateDocker_SingleServiceMode(t *testing.T) { + dir := t.TempDir() + path := writeTestFile(t, dir, "multi.yaml", fullMultiAgentYAML) + outputDir := filepath.Join(dir, "deploy") + + require.NoError(t, runGenerateDocker(&bytes.Buffer{}, []string{path}, outputDir, false)) + + compose, err := os.ReadFile(filepath.Join(outputDir, "docker-compose.yml")) + require.NoError(t, err) + assert.Contains(t, string(compose), "multi-agent:") + assert.NotContains(t, string(compose), "alpha:", "--multi-service=false must produce one service") +} + +func TestMultiAgentGenerateDocker_ReportsAMissingConfig(t *testing.T) { + dir := t.TempDir() + err := runGenerateDocker(&bytes.Buffer{}, []string{filepath.Join(dir, "absent.yaml")}, filepath.Join(dir, "deploy"), true) + require.Error(t, err) + assert.NoDirExists(t, filepath.Join(dir, "deploy")) +} + +// Regression: runGenerateK8s printed "Generating Kubernetes manifests..." and +// generated nothing, ignoring --output and --namespace entirely. +func TestMultiAgentGenerateK8s_WritesManifestsInTheNamespace(t *testing.T) { + dir := t.TempDir() + path := writeTestFile(t, dir, "multi.yaml", fullMultiAgentYAML) + outputDir := filepath.Join(dir, "k8s") + + var out bytes.Buffer + require.NoError(t, runGenerateK8s(&out, []string{path}, outputDir, "prod")) + + for _, name := range []string{"namespace.yaml", "deployment.yaml", "service.yaml"} { + content, err := os.ReadFile(filepath.Join(outputDir, name)) + require.NoError(t, err, "%s must be generated", name) + assert.Contains(t, string(content), "prod", "%s must use the requested namespace", name) + } + + deployment, err := os.ReadFile(filepath.Join(outputDir, "deployment.yaml")) + require.NoError(t, err) + assert.Contains(t, string(deployment), "replicas: 3", "the configured replica count must be used") +} + +// Regression: `--watch` fell into `select {}` and blocked forever after +// claiming to be watching for status changes. +func TestMultiAgentStatus_WatchIsReportedAsNotImplemented(t *testing.T) { + path := writeTestFile(t, t.TempDir(), "multi.yaml", fullMultiAgentYAML) + + done := make(chan error, 1) + go func() { done <- runMultiAgentStatus(&bytes.Buffer{}, []string{path}, "table", true) }() + + select { + case err := <-done: + require.Error(t, err) + assert.ErrorIs(t, err, errNotImplemented) + case <-time.After(10 * time.Second): + t.Fatal("multi-agent status --watch blocked instead of reporting that it is not implemented") + } +} + +func TestMultiAgentStatus_Formats(t *testing.T) { + path := writeTestFile(t, t.TempDir(), "multi.yaml", fullMultiAgentYAML) + + t.Run("table", func(t *testing.T) { + var out bytes.Buffer + require.NoError(t, runMultiAgentStatus(&out, []string{path}, "table", false)) + assert.Contains(t, out.String(), "alpha") + assert.Contains(t, out.String(), "does not contact a running deployment", + "the command reads the configuration; it must not imply live status") + }) + + t.Run("json", func(t *testing.T) { + var out bytes.Buffer + require.NoError(t, runMultiAgentStatus(&out, []string{path}, "json", false)) + + var decoded struct { + Agents map[string]struct { + Name string `json:"name"` + } `json:"agents"` + } + start := strings.Index(out.String(), "{") + require.GreaterOrEqual(t, start, 0) + require.NoError(t, json.Unmarshal([]byte(out.String()[start:]), &decoded)) + assert.Equal(t, "Alpha", decoded.Agents["alpha"].Name) + }) + + t.Run("unsupported", func(t *testing.T) { + err := runMultiAgentStatus(&bytes.Buffer{}, []string{path}, "toml", false) + require.Error(t, err) + }) +} + +func TestMultiAgentInit_CreatesAProjectItsOwnValidatorAccepts(t *testing.T) { + dir := chdirTemp(t) + + var out bytes.Buffer + require.NoError(t, runMultiAgentInit(&out, []string{"fleet"}, "basic", 2, "yaml", "path")) + + configPath := filepath.Join(dir, "fleet", "configs", "multi-agent.yaml") + assert.FileExists(t, configPath) + assert.FileExists(t, filepath.Join(dir, "fleet", "docker-compose.yml")) + assert.FileExists(t, filepath.Join(dir, "fleet", "k8s", "deployment.yaml")) + assert.FileExists(t, filepath.Join(dir, "fleet", "k8s", "service.yaml")) + assert.FileExists(t, filepath.Join(dir, "fleet", "README.md")) + assert.FileExists(t, filepath.Join(dir, "fleet", "agents", "agent-1", "config.yaml")) + assert.FileExists(t, filepath.Join(dir, "fleet", "agents", "agent-2", "config.yaml")) + + require.NoError(t, runMultiAgentValidate(&bytes.Buffer{}, []string{configPath}, false, true), + "the project init generates must pass the validate step init tells you to run") +} + +func TestMultiAgentInit_JSONFormat(t *testing.T) { + dir := chdirTemp(t) + + require.NoError(t, runMultiAgentInit(&bytes.Buffer{}, []string{"fleet"}, "basic", 1, "json", "path")) + + raw, err := os.ReadFile(filepath.Join(dir, "fleet", "configs", "multi-agent.json")) + require.NoError(t, err) + var decoded map[string]interface{} + require.NoError(t, json.Unmarshal(raw, &decoded)) + assert.Contains(t, decoded, "agents") +} + +func TestMultiAgentInit_RejectsBadInput(t *testing.T) { + dir := chdirTemp(t) + + for name, run := range map[string]func() error{ + "escaping name": func() error { + return runMultiAgentInit(&bytes.Buffer{}, []string{"../out"}, "basic", 1, "yaml", "path") + }, + "unknown format": func() error { return runMultiAgentInit(&bytes.Buffer{}, []string{"p"}, "basic", 1, "toml", "path") }, + "unknown template": func() error { + return runMultiAgentInit(&bytes.Buffer{}, []string{"p"}, "nonsense", 1, "yaml", "path") + }, + "unknown routing": func() error { return runMultiAgentInit(&bytes.Buffer{}, []string{"p"}, "basic", 1, "yaml", "carrier") }, + "zero agents": func() error { return runMultiAgentInit(&bytes.Buffer{}, []string{"p"}, "basic", 0, "yaml", "path") }, + } { + t.Run(name, func(t *testing.T) { + require.Error(t, run()) + }) + } + + assert.NoDirExists(t, filepath.Join(filepath.Dir(dir), "out")) +} + +func TestMultiAgentInit_EveryTemplateProducesAValidConfig(t *testing.T) { + for _, template := range []string{"basic", "microservices", "rag", "workflow"} { + t.Run(template, func(t *testing.T) { + dir := chdirTemp(t) + require.NoError(t, runMultiAgentInit(&bytes.Buffer{}, []string{"fleet"}, template, 3, "yaml", "path")) + require.NoError(t, runMultiAgentValidate(&bytes.Buffer{}, []string{ + filepath.Join(dir, "fleet", "configs", "multi-agent.yaml"), + }, false, true)) + }) + } +} + +// Regression: directory loading printed "Directory-based loading not yet +// implemented" and then returned nil, so a script checking the exit status was +// told the load had succeeded. +func TestMultiAgentLoad_DirectoryLoadingReportsNotImplemented(t *testing.T) { + dir := t.TempDir() + + err := runMultiAgentLoad(loadCommand(t, &bytes.Buffer{}), []string{dir}) + + require.Error(t, err) + assert.ErrorIs(t, err, errNotImplemented) +} + +func TestMultiAgentLoad_MissingPluginIsReported(t *testing.T) { + err := runMultiAgentLoad(loadCommand(t, &bytes.Buffer{}), []string{filepath.Join(t.TempDir(), "absent.so")}) + + require.Error(t, err) + assert.Contains(t, err.Error(), "failed to load plugin") +} + +func TestMultiAgentServe_RejectsBadInput(t *testing.T) { + dir := t.TempDir() + path := writeTestFile(t, dir, "multi.yaml", fullMultiAgentYAML) + + require.Error(t, runMultiAgentServe(context.Background(), &bytes.Buffer{}, []string{path}, "127.0.0.1", 0)) + require.Error(t, runMultiAgentServe(context.Background(), &bytes.Buffer{}, + []string{filepath.Join(dir, "absent.yaml")}, "127.0.0.1", 8080)) +} + +// The multi-agent server used to build a MultiAgentManager and then serve an +// unrelated, empty server while advertising per-agent endpoints. It now serves +// the agents from the configuration, so a bind failure must be reported rather +// than announced as a running server. +func TestMultiAgentServe_ReportsABindFailure(t *testing.T) { + path := writeTestFile(t, t.TempDir(), "multi.yaml", uniqueMultiAgentYAML(t)) + + var out bytes.Buffer + err := runMultiAgentServe(context.Background(), &out, []string{path}, "127.0.0.1", occupiedPort(t)) + + require.Error(t, err) + assert.Contains(t, err.Error(), "cannot bind") + assert.Contains(t, out.String(), "Registered 2 agent(s)") + assert.NotContains(t, out.String(), "listening") +} + +// uniqueMultiAgentYAML returns a configuration whose agent IDs are unique to +// the test: agent definitions are registered in a process-wide registry, and +// re-registering the same ID fails. +func uniqueMultiAgentYAML(t *testing.T) string { + t.Helper() + + prefix := strings.ToLower(strings.ReplaceAll(t.Name(), "/", "-")) + return strings.NewReplacer( + "alpha", prefix+"-alpha", + "beta", prefix+"-beta", + ).Replace(fullMultiAgentYAML) +} + +// loadCommand returns the registered "multi-agent load" command with its +// output redirected for the duration of the test. +func loadCommand(t *testing.T, out *bytes.Buffer) *cobra.Command { + t.Helper() + + for _, sub := range multiAgentCmd.Commands() { + if sub.Name() != "load" { + continue + } + sub.SetOut(out) + t.Cleanup(func() { sub.SetOut(nil) }) + return sub + } + t.Fatal("the multi-agent load command is not registered") + return nil +} diff --git a/docs/PRODUCTION.md b/docs/PRODUCTION.md new file mode 100644 index 0000000..6268c9c --- /dev/null +++ b/docs/PRODUCTION.md @@ -0,0 +1,388 @@ +# Running GoLangGraph in production + +This document covers the parts of GoLangGraph you need to configure +deliberately before exposing it to real traffic: authentication, cross-origin +access, tool sandboxing, durable execution and health checking. + +Defaults are chosen so that local development works out of the box. Several of +them are **not** the right choice for a deployment, and each is called out +below. + +--- + +## 1. Server security + +`ServerConfig.Security` controls authentication, allowed origins and request +limits. It defaults to `DefaultSecurityConfig()`, which is permissive. + +```go +cfg := server.DefaultServerConfig() +cfg.Security = &server.SecurityConfig{ + RequireAuth: true, + APIKeys: []string{os.Getenv("GOLANGGRAPH_API_KEY")}, + AllowedOrigins: []string{"https://studio.example.com"}, + MaxRequestBytes: 4 << 20, + PublicPaths: []string{"/api/v1/health"}, +} +srv := server.NewServer(cfg) +``` + +| Setting | Default | What to set in production | +| --- | --- | --- | +| `RequireAuth` | `false` | `true`. With it off, every endpoint is unauthenticated. | +| `APIKeys` | empty | At least one key. Clients send it as `X-API-Key`. Keys are compared in constant time. | +| `AllowedOrigins` | empty (any) | The exact origins that may call the API. This list also governs **WebSocket** upgrades. | +| `MaxRequestBytes` | 4 MiB | Lower it if your payloads are small. | +| `PublicPaths` | `/api/v1/health` | Paths that skip authentication, for load-balancer probes. | + +### Both servers, not just one + +There are two serving paths in this project and **each needs configuring**: + +| Type | Constructor | Used by | +| --- | --- | --- | +| `server.Server` | `server.NewServer(cfg)` | The `serve` command and the `/api/v1` API that Studio talks to | +| `server.AutoServer` | `server.NewAutoServer(cfg)` | The `auto-serve` command and auto-generated per-agent endpoints | + +`AutoServerConfig` takes the same `Security` value: + +```go +cfg := server.DefaultAutoServerConfig() +cfg.Security = &server.SecurityConfig{ + RequireAuth: true, + APIKeys: []string{os.Getenv("GOLANGGRAPH_API_KEY")}, + AllowedOrigins: []string{"https://studio.example.com"}, + PublicPaths: []string{"/health"}, +} +cfg.MaxRequestSize = 4 << 20 +``` + +On `AutoServer`, panic recovery, the request size limit and the security headers +are unconditional. They were previously opt-in through the `Middleware` name +list, so a deployment that left `"recovery"` out of that list had a panicking +handler tear down the connection. + +**Agent registry isolation.** `NewAutoServer` serves from the process-wide agent +registry, so two auto-servers in one process see each other's agents β€” which is +rarely what you want if they differ in exposure or credentials. Use +`server.NewAutoServerWithRegistry(cfg, agent.NewAgentRegistry())` to give a +server its own. Endpoints also cannot be regenerated once `Start` has been +called; register agents before starting. + +**Why the origin list matters for WebSockets.** A WebSocket upgrade that accepts +any origin allows cross-site WebSocket hijacking: any page a signed-in user +visits can open a socket to your server and drive it as that user. Setting +`AllowedOrigins` closes that. Leaving it empty accepts any origin and is only +appropriate for local development. + +Authentication fails closed: if `RequireAuth` is true and no keys are +configured, every request is rejected rather than allowed. + +### Always-on protections + +These apply regardless of configuration: + +- Request bodies are size-limited. +- Handler panics become a 500 with a generic JSON body; the stack is logged, never sent to the client. +- `X-Content-Type-Options`, `X-Frame-Options` and `Referrer-Policy` are set. +- Cross-origin preflight is answered for every route. + +--- + +## 2. Tool sandboxing + +Tool arguments are produced by a language model, which may be steered by +untrusted input. Treat every built-in tool as attacker-controlled and bound it +with a `tools.SecurityPolicy`. + +```go +policy := tools.DefaultSecurityPolicy() +policy.SetAllowedRoots([]string{"/srv/agent-workspace"}) +policy.MaxOutputBytes = 256 << 10 +policy.AllowedCommands = []string{"ls", "wc"} +policy.AllowedHosts = []string{"api.example.com"} + +read := tools.NewFileReadTool() +read.SetSecurityPolicy(policy) +``` + +**Filesystem tools** are confined to `AllowedRoots` (working directory and temp +directory by default). Paths are resolved and symlinks followed during +validation, so a link planted inside an allowed root cannot reach outside one. +An extension allowlist alone is not a boundary: `.json` and `.yaml` files +elsewhere on a host include kubeconfigs, registry credentials and cloud +credential files. + +**The shell tool** runs an allowlisted command directly, never through a shell. +Note what the allowlist cannot contain: `find` executes arbitrary programs via +`-exec`, and `cat`/`grep` read any file the process can reach, so none of them +are in the default set. Arguments containing program-executing flags or shell +metacharacters are rejected, as is a command given as a path. + +**The HTTP tool** refuses loopback, private, link-local, multicast and +unspecified addresses unless `AllowPrivateNetwork` is set. This blocks +server-side request forgery, including the cloud instance metadata endpoint at +`169.254.169.254`. The check runs at dial time against the resolved address, so +a hostname that resolves inward is caught even if it resolved elsewhere a +moment earlier, and every redirect hop is re-validated. + +If your agent legitimately needs an internal service, prefer +`AllowedHosts` over `AllowPrivateNetwork`. + +--- + +## 2b. Multi-agent deployments + +`MultiAgentManager` serves its own HTTP surface, and several of its controls +were previously inert. Configure them deliberately: + +- **Authentication** accepted any non-empty API key. It now compares against + configured keys in constant time, and enabling auth without configuring keys + fails construction rather than admitting everyone. +- **`GET /config`** returned provider API keys, database, cache and SMTP + passwords, the Slack webhook and both secret maps verbatim to any caller. + The response is redacted; treat any deployment that exposed it as having + leaked those credentials. +- **Rate limiting** was configured and did nothing. Global, per-IP and + per-agent limits are enforced and return 429 with `Retry-After`. `per_user` + is not enforceable without request identity and is logged loudly at startup + rather than silently ignored. +- **`GET /health`** always answered 200 "healthy", so every liveness probe + passed regardless of agent state. It now aggregates and returns 503. +- **Routing rule conditions** were parsed, stored and never compared against a + request, so a rule guarded by a condition matched everything. +- **Health checkers** leaked a goroutine per agent that nothing could stop. + They now start with `Start` and are cancelled and joined by `Stop`. + +Config validation is stricter as a result: patternless rules, duplicate rule +IDs, invalid regexes, negative timeouts, inverted scaling bounds and routes to +unknown or disabled agents are now rejected at load. A configuration that was +previously accepted while being partly inert may now fail to start β€” that is +the point. + +## 3. Durable execution and resume + +Attach a checkpointer to persist state after every node, so a run that dies +mid-graph can resume instead of starting over. + +```go +checkpointer := persistence.NewFileCheckpointer("/var/lib/golanggraph/checkpoints") +saver := persistence.NewCheckpointSaver(checkpointer) + +graph.WithCheckpointer(saver, threadID) +``` + +To resume after a crash, load the most recent checkpoint and restart from the +node that should run next. `latest.NodeID` is the node that *completed*, so +resume from its successor: + +```go +latest, err := persistence.Latest(ctx, checkpointer, threadID) +if err != nil { + return err +} +if latest != nil { + next, err := graph.GetNextNodes(ctx, latest.NodeID, latest.State) + if err != nil { + return err + } + if len(next) > 0 { + _, err = graph.ExecuteWithOptions(ctx, latest.State, &core.ExecuteOptions{ + ThreadID: threadID, + StartNode: next[0], + }) + } +} +``` + +Thread and checkpoint identifiers become path components in the file backend and +are validated; values containing separators or `..` are rejected. + +### Database backends + +- **PostgreSQL** could not be used as a `Checkpointer` at all: `checkpoints` + has a foreign key to `threads`, nothing in the interface created the thread + row, so the first save of any new thread failed on the constraint. The thread + is now upserted in the same transaction. +- **The RAG document path never worked.** `SaveDocument` and vector search + passed a Go map and `[]float64` straight to `database/sql`, which rejects + both, so every call failed. Embeddings read back were also discarded. +- **Redis leaked state between threads.** Keys were built by joining the thread + and checkpoint IDs with `:` without escaping, so thread `x:a` + checkpoint + `b:c1` collided with thread `x:a:b` + checkpoint `c1` β€” loading one thread + could return another's checkpoint. Thread IDs are commonly user- or + session-derived, making this a cross-tenant leak. Segments are now escaped + and the loaded thread ID is verified. +- **Redis expiry was hardcoded to 24 hours**, so every deployment silently lost + its checkpoints after a day. Set `DatabaseConfig.CheckpointTTL`. +- `rows.Err()` was never checked when listing, so a connection dropping + mid-iteration returned a silently truncated list with a nil error β€” and + `Latest()` is built on that list, so a resume could silently pick up from the + wrong checkpoint. + +Postgres and Redis now have real integration tests. They skip with an explicit +message when no server is reachable, and hard-fail if `POSTGRES_TEST_DSN` or +`REDIS_TEST_ADDR` is set but unreachable. + +### Human-in-the-loop + +```go +graph.Config.InterruptBefore = []string{"apply_changes"} + +_, err := graph.Execute(ctx, state) + +var interrupt *core.InterruptError +if errors.As(err, &interrupt) { + interrupt.State.Set("amount", reviewedAmount) // a person edits the state + graph.Config.InterruptBefore = nil + final, err := graph.Resume(ctx, interrupt) +} +``` + +An interrupt is a normal, resumable outcome, not a failure. Over HTTP it is +reported as `200` with `"status": "interrupted"`. + +--- + +## 4. Health checking + +`golanggraph health` answers two different questions; pick the right one. + +```bash +# Is this server serving? Use this as a container health check. +golanggraph health --server http://127.0.0.1:8080 + +# Are this host's configured dependencies reachable? +POSTGRES_HOST=db REDIS_HOST=cache golanggraph health +``` + +The dependency scan probes only the services that are actually configured +(`POSTGRES_HOST`, `REDIS_HOST`, `OLLAMA_URL`), because defaulting to localhost +would report a failure in every deployment that does not use them. Missing +optional provider credentials are warnings and exit `0`; pass `--strict` to make +warnings fail. + +Do not use the dependency scan as a container health check: an absent optional +dependency would restart a perfectly healthy container forever. + +--- + +## 5. Error handling + +Execution returns typed sentinels; branch with `errors.Is` rather than matching +message text. + +| Sentinel | Meaning | +| --- | --- | +| `core.ErrGraphInvalid` | The graph failed validation. | +| `core.ErrRecursionLimit` | `MaxIterations` was exceeded (LangGraph's `GraphRecursionError`). | +| `core.ErrInterrupted` | Paused at an interrupt, or stopped by `Interrupt()`. | +| `core.ErrNodePanic` | A node or condition panicked; recovered and converted. | +| `core.ErrNoRoute` | No outgoing edge matched. | +| `core.ErrGraphClosed` | The graph was closed. | +| `llm.ErrProviderUnavailable` | Transient provider failure; retrying may help. | +| `llm.ErrRateLimited` | The provider asked you to slow down. | +| `llm.ErrProviderAuth` | Credentials were rejected. | +| `llm.ErrProviderRequest` | Permanent client error; retrying will not help. | + +The original cause is always wrapped, so `errors.Is` against your own sentinel +works through the engine. + +On failure `Execute` returns the last known good state alongside the error, so +partial progress is inspectable without a checkpointer. + +--- + +## 6. Retries + +Node retries are **off by default**. Node bodies commonly perform +non-idempotent work β€” model calls, tool side effects, writes β€” so retrying them +silently can duplicate that work. + +Enable them where the work is safe to repeat: + +```go +node := graph.AddNode("fetch", "Fetch", fetchFn) +node.Retry = &core.RetryPolicy{ + MaxAttempts: 3, + Delay: time.Second, + Backoff: 2, + RetryIf: func(err error) bool { return errors.Is(err, llm.ErrProviderUnavailable) }, +} +``` + +Each attempt starts from the state as it was before the attempt, so a partially +mutated state from a failed try cannot leak into the retry. + +Provider-level retries are separate and driven by `ProviderConfig.RetryCount` / +`RetryDelay`. They apply only to transient failures and honour a `Retry-After` +header. + +--- + +## 7. Concurrency + +- A `*core.Graph` is safe to execute concurrently; each run keeps its own state and history. +- `GetCurrentState()` and `GetExecutionHistory()` reflect the **most recent** run and are for observability. The authoritative result of a run is what `Execute` returns. +- An `*agent.Agent` rejects a second concurrent run, because its conversation and execution record are per-agent. +- `Graph.Stream()` is shared and lossy: results are dropped rather than stalling execution. For lossless streaming, pass a channel via `ExecuteOptions.Stream`, which receives only that run's steps and is closed when the run ends. + +--- + +## 8. Observability + +Every node execution produces an `ExecutionResult` carrying the node ID, step +index, success, duration, attempt count and the state after the node. Failures +are recorded too, with a serialisable `error` field. + +Agent executions record `execution_path` (the nodes that ran) and +`state_changes` (a before/after snapshot per node), which is what a debugging +client such as GoLangGraph Studio renders. + +--- + +## Behaviour changes to be aware of + +If you are upgrading, these defaults and shapes changed: + +### `AgentExecution` wire format + +`agent.AgentExecution` carried no struct tags, so it alone on this API +serialised with Go's default PascalCase field names while every neighbouring +payload (agent configs, providers, graph topology) was snake_case. Worse, its +`Error` field is a Go `error`, which `encoding/json` renders as `{}` β€” a failed +execution reached the client with no reason in it at all. + +Every field is now tagged. The wire format is snake_case, and the failure +reason travels as a string `error` field; the Go `error` stays on the struct +for in-process callers and is excluded from JSON. `pkg/server`'s +`TestFrontendAPIContract` pins both the tagged names and the absence of the +untagged ones. + +GoLangGraph Studio consumes this shape and is updated in lockstep. Any other +client reading `ID`/`Input`/`Output`/`Success` must move to +`id`/`input`/`output`/`success`. + +| Change | Before | Now | +| --- | --- | --- | +| Node retries | 3 attempts by default | Off by default; opt in per node | +| `AgentExecution.Error` | A Go `error`, serialised as `{}` | String `error` field carrying the reason | +| `GET /api/v1/agents` | List of ID strings | List of agent configurations | +| `GET /api/v1/providers` | List of name strings | List of provider descriptions, credentials omitted | +| `BaseState` JSON | Serialised as `{}` | Full `{"data":…,"metadata":…}` payload | +| Conditional edges | Recorded but never used during execution | Evaluated once per visit and routed | +| Container health check | Local dependency scan | Server endpoint probe | +| `AutoServer` auth | None available | Configurable via `Security`, off by default | +| `AutoServer` CORS | Hardcoded `*` | Configurable allowlist | +| `AutoServer` `MaxRequestSize` | Declared, never enforced | Enforced | +| `AutoServer.GenerateEndpoints` | Callable at any time | Refused once `Start` has run | +| Agent IDs | Empty when built from a config literal | Assigned automatically by `NewAgent` | +| `builder` provider fallback | Returned `"mock"`, which does not exist | Returns empty and warns | +| `AgentSwarm.Execute` result order | Map iteration order (random) | The order agents were supplied | +| `ServerConfig.LogLevel` | Declared, never read | Applied to the server logger | +| `AgentExecution` JSON | Untagged, so Go PascalCase (`ID`, `Input`, …) | Tagged snake_case (`id`, `input`, …) like the rest of the API | +| Agent turn after the first | Treated as an interrupt resume; the new input was dropped | A new turn; resume is set only by `SeedResumeState`/`SeedConversation` | +| `multi-agent init` | Panicked writing project YAML (`cannot marshal type: llm.EarlyExitFunc`) | Writes the project; func-typed config fields are `yaml:"-"` | + +See `test/conformance/DEVIATIONS.md` for where GoLangGraph intentionally differs +from LangGraph, and why. diff --git a/examples/01-basic-chat/01-basic-chat b/examples/01-basic-chat/01-basic-chat deleted file mode 100755 index 21911ea..0000000 Binary files a/examples/01-basic-chat/01-basic-chat and /dev/null differ diff --git a/examples/01-basic-chat/basic-chat b/examples/01-basic-chat/basic-chat deleted file mode 100755 index 21911ea..0000000 Binary files a/examples/01-basic-chat/basic-chat and /dev/null differ diff --git a/examples/02-react-agent/02-react-agent b/examples/02-react-agent/02-react-agent deleted file mode 100755 index 17ad331..0000000 Binary files a/examples/02-react-agent/02-react-agent and /dev/null differ diff --git a/examples/02-react-agent/react-agent b/examples/02-react-agent/react-agent deleted file mode 100755 index 17ad331..0000000 Binary files a/examples/02-react-agent/react-agent and /dev/null differ diff --git a/examples/03-multi-agent/03-multi-agent b/examples/03-multi-agent/03-multi-agent deleted file mode 100755 index 0fad0ba..0000000 Binary files a/examples/03-multi-agent/03-multi-agent and /dev/null differ diff --git a/examples/04-rag-system/04-rag-system b/examples/04-rag-system/04-rag-system deleted file mode 100755 index ee11b24..0000000 Binary files a/examples/04-rag-system/04-rag-system and /dev/null differ diff --git a/examples/05-streaming/05-streaming b/examples/05-streaming/05-streaming deleted file mode 100755 index f56d308..0000000 Binary files a/examples/05-streaming/05-streaming and /dev/null differ diff --git a/examples/06-persistence/06-persistence b/examples/06-persistence/06-persistence deleted file mode 100755 index f74528b..0000000 Binary files a/examples/06-persistence/06-persistence and /dev/null differ diff --git a/examples/07-tools-integration/go.sum b/examples/07-tools-integration/go.sum index e2dd872..8802d11 100644 --- a/examples/07-tools-integration/go.sum +++ b/examples/07-tools-integration/go.sum @@ -1,4 +1,5 @@ github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc h1:U9qPSI2PIWSS1VwoXQT9A3Wy9MM3WgvqSxFWenqJduM= +github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/go-ole/go-ole v1.2.6 h1:/Fpf6oFPoeFik9ty7siob0G6Ke8QvQEuVcuChpwXzpY= github.com/go-ole/go-ole v1.2.6/go.mod h1:pprOEPIfldk/42T2oK7lQ4v4JSDwmV0As9GaiUsvbm0= github.com/google/go-cmp v0.5.6/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE= @@ -7,6 +8,7 @@ github.com/google/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeN github.com/lufia/plan9stats v0.0.0-20211012122336-39d0f177ccd0 h1:6E+4a0GO5zZEnZ81pIr0yLvtUWk2if982qA3F3QD6H4= github.com/lufia/plan9stats v0.0.0-20211012122336-39d0f177ccd0/go.mod h1:zJYVVT2jmtg6P3p1VtQj7WsuWi/y4VnjVBn7F8KPB3I= github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 h1:Jamvg5psRIccs7FGNTlIRMkT8wgtp5eCXdBlqhYGL6U= +github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= github.com/power-devops/perfstat v0.0.0-20210106213030-5aafc221ea8c h1:ncq/mPwQF4JjgDlrVEn3C11VoGHZN7m8qihwgMEtzYw= github.com/power-devops/perfstat v0.0.0-20210106213030-5aafc221ea8c/go.mod h1:OmDBASR4679mdNQnz2pUhc2G8CO2JrUAVFDRBDP/hJE= github.com/shirou/gopsutil/v3 v3.24.5 h1:i0t8kL+kQTvpAYToeuiVk3TgDeKOFioZO3Ztz/iZ9pI= @@ -16,6 +18,7 @@ github.com/shoenig/go-m1cpu v0.1.6/go.mod h1:1JJMcUBvfNwpq05QDQVAnx3gUHr9IYF7GNg github.com/shoenig/test v0.6.4 h1:kVTaSd7WLz5WZ2IaoM0RSzRsUD+m8wRR+5qvntpn4LU= github.com/shoenig/test v0.6.4/go.mod h1:byHiCGXqrVaflBLAMq/srcZIHynQPQgeyvkvXnjqq0k= github.com/stretchr/testify v1.10.0 h1:Xv5erBjTwe/5IxqUQTdXv5kgmIvbHo3QQyRwhJsOfJA= +github.com/stretchr/testify v1.10.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY= github.com/tklauser/go-sysconf v0.3.12 h1:0QaGUFOdQaIVdPgfITYzaTegZvdCjmYO52cSFAEVmqU= github.com/tklauser/go-sysconf v0.3.12/go.mod h1:Ho14jnntGE1fpdOqQEEaiKRpvIavV0hSfmBq8nJbHYI= github.com/tklauser/numcpus v0.6.1 h1:ng9scYS7az0Bk4OZLvrNXNSAO2Pxr1XXRAPyjhIx+Fk= @@ -27,6 +30,7 @@ golang.org/x/sys v0.0.0-20201204225414-ed752295db88/go.mod h1:h1NjWce9XRLGQEsW7w golang.org/x/sys v0.8.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.11.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.33.0 h1:q3i8TbbEz+JRD9ywIRlyRAQbM0qF7hu24q3teo2hbuw= +golang.org/x/sys v0.33.0/go.mod h1:BJP2sWEmIv4KK5OTEluFJCKSidICx8ciO85XgH3Ak8k= golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= diff --git a/examples/07-tools-integration/tools-integration b/examples/07-tools-integration/tools-integration deleted file mode 100755 index cbd89d6..0000000 Binary files a/examples/07-tools-integration/tools-integration and /dev/null differ diff --git a/examples/08-production-ready/go.sum b/examples/08-production-ready/go.sum index 5f91d64..feca7ad 100644 --- a/examples/08-production-ready/go.sum +++ b/examples/08-production-ready/go.sum @@ -9,7 +9,9 @@ github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc/go.mod h1:J7Y8Yc github.com/frankban/quicktest v1.14.6 h1:7Xjx+VpznH+oBnejlPUj8oUpdxnVs4f8XU8WnHkI4W8= github.com/frankban/quicktest v1.14.6/go.mod h1:4ptaffx2x8+WTWXmUCuVU6aPUX1/Mz7zb5vbUoiM6w0= github.com/fsnotify/fsnotify v1.8.0 h1:dAwr6QBTBZIkG8roQaJjGof0pp0EeF+tNV7YBP3F/8M= +github.com/fsnotify/fsnotify v1.8.0/go.mod h1:8jBTzvmWwFyi3Pb8djgCCO5IBqzKJ/Jwo8TRcHyHii0= github.com/go-viper/mapstructure/v2 v2.2.1 h1:ZAaOCxANMuZx5RCeg0mBdEZk7DZasvvZIxtHqx8aGss= +github.com/go-viper/mapstructure/v2 v2.2.1/go.mod h1:oJDH3BJKyqBA2TXFhDsKDGDTlndYOZ6rGS0BRZIxGhM= github.com/google/go-cmp v0.6.0 h1:ofyhxvXcZhMsU5ulbFiLKl/XBFqE1GSq7atu8tAmTRI= github.com/google/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY= github.com/gorilla/mux v1.8.1 h1:TuBL49tXwgrFYWhqrNgrUNEY92u81SPhu7sTdzQEiWY= @@ -27,6 +29,7 @@ github.com/kylelemons/godebug v1.1.0/go.mod h1:9/0rRGxNHcop5bhtWyNeEfOS8JIWk580+ github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 h1:C3w9PqII01/Oq1c1nUAm88MOHcQC9l5mIlSMApZMrHA= github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822/go.mod h1:+n7T8mK8HuQTcFwEeznm/DIxMOiR9yIdICNftLE1DvQ= github.com/pelletier/go-toml/v2 v2.2.3 h1:YmeHyLY8mFWbdkNWwpr+qIL2bEqT0o95WSdkNHvL12M= +github.com/pelletier/go-toml/v2 v2.2.3/go.mod h1:MfCQTFTvCcUyyvvwm1+G6H/jORL20Xlb6rzQu9GuUkc= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 h1:Jamvg5psRIccs7FGNTlIRMkT8wgtp5eCXdBlqhYGL6U= github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= @@ -41,18 +44,24 @@ github.com/prometheus/procfs v0.15.1/go.mod h1:fB45yRUv8NstnjriLhBQLuOUt+WW4BsoG github.com/rogpeppe/go-internal v1.10.0 h1:TMyTOH3F/DB16zRVcYyreMH6GnZZrwQVAoYjRBZyWFQ= github.com/rogpeppe/go-internal v1.10.0/go.mod h1:UQnix2H7Ngw/k4C5ijL5+65zddjncjaFoBhdsK/akog= github.com/sagikazarmark/locafero v0.7.0 h1:5MqpDsTGNDhY8sGp0Aowyf0qKsPrhewaLSsFaodPcyo= +github.com/sagikazarmark/locafero v0.7.0/go.mod h1:2za3Cg5rMaTMoG/2Ulr9AwtFaIppKXTRYnozin4aB5k= github.com/sirupsen/logrus v1.9.3 h1:dueUQJ1C2q9oE3F7wvmSGAaVtTmUizReu6fjN8uqzbQ= github.com/sirupsen/logrus v1.9.3/go.mod h1:naHLuLoDiP4jHNo9R0sCBMtWGeIprob74mVsIT4qYEQ= github.com/sourcegraph/conc v0.3.0 h1:OQTbbt6P72L20UqAkXXuLOj79LfEanQ+YQFNpLA9ySo= github.com/sourcegraph/conc v0.3.0/go.mod h1:Sdozi7LEKbFPqYX2/J+iBAM6HpqSLTASQIKqDmF7Mt0= github.com/spf13/afero v1.12.0 h1:UcOPyRBYczmFn6yvphxkn9ZEOY65cpwGKb5mL36mrqs= +github.com/spf13/afero v1.12.0/go.mod h1:ZTlWwG4/ahT8W7T0WQ5uYmjI9duaLQGy3Q2OAl4sk/4= github.com/spf13/cast v1.7.1 h1:cuNEagBQEHWN1FnbGEjCXL2szYEXqfJPbP2HNUaca9Y= +github.com/spf13/cast v1.7.1/go.mod h1:ancEpBxwJDODSW/UG4rDrAqiKolqNNh2DX3mk86cAdo= github.com/spf13/pflag v1.0.6 h1:jFzHGLGAlb3ruxLB8MhbI6A8+AQX/2eW4qeyNZXNp2o= +github.com/spf13/pflag v1.0.6/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg= github.com/spf13/viper v1.20.1 h1:ZMi+z/lvLyPSCoNtFCpqjy0S4kPbirhpTMwl8BkW9X4= +github.com/spf13/viper v1.20.1/go.mod h1:P9Mdzt1zoHIG8m2eZQinpiBjo6kCmZSKBClNNqjJvu4= github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI= github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= github.com/stretchr/testify v1.10.0 h1:Xv5erBjTwe/5IxqUQTdXv5kgmIvbHo3QQyRwhJsOfJA= +github.com/stretchr/testify v1.10.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY= github.com/subosito/gotenv v1.6.0 h1:9NlTDc1FTs4qu0DDq7AEtTPNw6SVm7uBMsUCUjABIf8= github.com/subosito/gotenv v1.6.0/go.mod h1:Dk4QP5c2W3ibzajGcXpNraDfq2IrhjMIvMSWPKKo0FU= go.uber.org/atomic v1.9.0 h1:ECmE8Bn/WFTYwEW/bpKD3M8VtR/zQVbavAoalC1PYyE= @@ -61,8 +70,11 @@ go.uber.org/multierr v1.9.0 h1:7fIwc/ZtS0q++VgcfqFDxSBZVv/Xo49/SYnDFupUwlI= go.uber.org/multierr v1.9.0/go.mod h1:X2jQV1h+kxSjClGpnseKVIxpmcjrj7MNnI0bnlfKTVQ= golang.org/x/sys v0.0.0-20220715151400-c0bba94af5f8/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.33.0 h1:q3i8TbbEz+JRD9ywIRlyRAQbM0qF7hu24q3teo2hbuw= +golang.org/x/sys v0.33.0/go.mod h1:BJP2sWEmIv4KK5OTEluFJCKSidICx8ciO85XgH3Ak8k= golang.org/x/text v0.26.0 h1:P42AVeLghgTYr4+xUnTRKDMqpar+PtX7KWuNQL21L8M= +golang.org/x/text v0.26.0/go.mod h1:QK15LZJUUQVJxhz7wXgxSy/CJaTFjd0G+YLonydOVQA= google.golang.org/protobuf v1.36.1 h1:yBPeRvTftaleIgM3PZ/WBIZ7XM/eEYAaEyCwvyjq/gk= +google.golang.org/protobuf v1.36.1/go.mod h1:9fA7Ob0pmnwhb644+1+CVWFRbNajQ6iRojtC/QF5bRE= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk= gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q= diff --git a/examples/08-production-ready/production-ready b/examples/08-production-ready/production-ready deleted file mode 100755 index 0d1f01a..0000000 Binary files a/examples/08-production-ready/production-ready and /dev/null differ diff --git a/examples/09-workflow-graph/workflow-graph b/examples/09-workflow-graph/workflow-graph deleted file mode 100755 index e5aa878..0000000 Binary files a/examples/09-workflow-graph/workflow-graph and /dev/null differ diff --git a/pkg/agent/agent.go b/pkg/agent/agent.go index b7c02e3..9135ebd 100644 --- a/pkg/agent/agent.go +++ b/pkg/agent/agent.go @@ -47,7 +47,7 @@ type AgentConfig struct { StreamingMode llm.StreamMode `json:"streaming_mode,omitempty"` // EarlyExit cancels remaining stream tokens once a complete JSON/tool-call // is formed. Nil disables token-stream early-exit (multipass JSON exit still applies). - EarlyExit llm.EarlyExitFunc `json:"-"` + EarlyExit llm.EarlyExitFunc `json:"-" yaml:"-"` Timeout time.Duration `json:"timeout"` Metadata map[string]interface{} `json:"metadata"` Middleware []Middleware `json:"-"` @@ -202,15 +202,29 @@ type BaseAgent struct { currentIteration int executionHistory []AgentExecution pendingToolCalls []llm.ToolCall // seeded on HITL resume (mid-tool-call) + // resumeSeeded marks a run that continues an interrupted one, set only by + // the Seed* entry points. Resume used to be inferred from a non-empty + // conversation, which is also true of every ordinary turn after the first: + // the second question a caller asked was treated as a resume, so it was + // never added to the conversation and the agent answered the first one + // again. Resume is a property of how the run was set up, not of history. + resumeSeeded bool } -// ... (AgentExecution struct remains same) - // NewAgent creates a new base agent func NewAgent(config *AgentConfig, llmManager *llm.ProviderManager, toolRegistry *tools.ToolRegistry) *BaseAgent { // Create a copy of config to avoid modification of original agentConfig := *config + // An agent without an ID is indistinguishable from every other such agent: + // AgentManager keys its map by ID, so a second one silently replaces the + // first. Building AgentConfig as a literal β€” which the documented examples + // do β€” leaves ID empty, so assign one here rather than requiring every + // caller to remember. + if agentConfig.ID == "" { + agentConfig.ID = uuid.New().String() + } + // Validate and sanitize configuration if err := agentConfig.ValidateAndSanitize(); err != nil { // Log the error and apply default configuration @@ -326,8 +340,14 @@ func (a *BaseAgent) buildReActGraph() { a.graph.AddEdge("observe", "finalize", a.shouldContinueReasoning) // Set start and end nodes - a.graph.SetStartNode("reason") - a.graph.AddEndNode("finalize") + // A failure here means the graph is malformed; record it so Validate + // reports it instead of the agent running against a broken graph. + if err := a.graph.SetStartNode("reason"); err != nil { + a.logger.WithError(err).Error("failed to set graph start node") + } + if err := a.graph.AddEndNode("finalize"); err != nil { + a.logger.WithError(err).Error("failed to add graph end node") + } } // buildChatGraph builds a simple chat graph @@ -339,8 +359,14 @@ func (a *BaseAgent) buildChatGraph() { chatNode.Metadata["type"] = "chat" // Set start and end nodes - a.graph.SetStartNode("chat") - a.graph.AddEndNode("chat") + // A failure here means the graph is malformed; record it so Validate + // reports it instead of the agent running against a broken graph. + if err := a.graph.SetStartNode("chat"); err != nil { + a.logger.WithError(err).Error("failed to set graph start node") + } + if err := a.graph.AddEndNode("chat"); err != nil { + a.logger.WithError(err).Error("failed to add graph end node") + } } // buildToolGraph builds a tool-focused graph @@ -361,8 +387,25 @@ func (a *BaseAgent) buildToolGraph() { a.graph.AddEdge("review", "plan", a.shouldReplan) // Set start and end nodes - a.graph.SetStartNode("plan") - a.graph.AddEndNode("review") + // A failure here means the graph is malformed; record it so Validate + // reports it instead of the agent running against a broken graph. + if err := a.graph.SetStartNode("plan"); err != nil { + a.logger.WithError(err).Error("failed to set graph start node") + } + if err := a.graph.AddEndNode("review"); err != nil { + a.logger.WithError(err).Error("failed to add graph end node") + } +} + +// nodeName resolves a node's display name for state-change records. +func (a *BaseAgent) nodeName(nodeID string) string { + if a.graph == nil { + return nodeID + } + if node, ok := a.graph.Nodes[nodeID]; ok && node != nil && node.Name != "" { + return node.Name + } + return nodeID } // Execute executes the agent with the given input @@ -383,10 +426,12 @@ func (a *BaseAgent) ExecuteThread(ctx context.Context, threadID string, input st resumeIter = 0 } // Fresh runs reset iteration; seeded resume keeps currentIteration. - if a.conversation.Size() == 0 { + resuming := a.resumeSeeded + if !resuming { a.currentIteration = 0 resumeIter = 0 } + a.resumeSeeded = false // consumed: the next run is fresh unless re-seeded a.mu.Unlock() defer func() { @@ -416,7 +461,6 @@ func (a *BaseAgent) ExecuteThread(ctx context.Context, threadID string, input st } execution.Input = input // Update input in execution record if modified - resuming := a.conversation.Size() > 0 if resuming { execution.Metadata["resumed"] = true execution.Metadata["resume_iteration"] = resumeIter @@ -450,10 +494,58 @@ func (a *BaseAgent) ExecuteThread(ctx context.Context, threadID string, input st // Let's assume 'state' is fresh for this turn. } - // Execute the graph - finalState, err := a.graph.Execute(ctx, state) + // Execute the graph, collecting the steps as they happen. + // + // ExecutionPath and StateChanges were declared but never populated, so a + // debugging client had no way to see which nodes ran: GoLangGraph Studio + // highlights nodes from execution_path, and an empty list means its graph + // view shows nothing for a run that did execute. + steps := make(chan *core.ExecutionResult, 256) + collected := make(chan struct { + path []string + changes []StateChange + }, 1) + + go func() { + var path []string + var changes []StateChange + var previous map[string]interface{} + + for result := range steps { + path = append(path, result.NodeID) + + change := StateChange{ + NodeID: result.NodeID, + NodeName: a.nodeName(result.NodeID), + Timestamp: result.Timestamp, + Before: previous, + } + if result.State != nil { + after := make(map[string]interface{}, len(result.State.GetAll())) + for k, v := range result.State.GetAll() { + after[k] = v + } + change.After = after + previous = after + } + changes = append(changes, change) + } + + collected <- struct { + path []string + changes []StateChange + }{path: path, changes: changes} + }() + + finalState, err := a.graph.ExecuteWithOptions(ctx, state, &core.ExecuteOptions{Stream: steps}) + + observed := <-collected + execution.ExecutionPath = observed.path + execution.StateChanges = observed.changes + if err != nil { execution.Error = err + execution.ErrorMessage = err.Error() execution.Success = false } else { execution.Success = true @@ -497,9 +589,10 @@ func (a *BaseAgent) ExecuteThread(ctx context.Context, threadID string, input st } } - // Track execution path from graph - if a.graph != nil { - // This would be populated by the graph + // ExecutionPath is normally collected from the graph's own step stream + // above. Fall back to a path a node published in state, for graphs that + // report their own route; appending unconditionally would double it. + if a.graph != nil && len(execution.ExecutionPath) == 0 { if executionPathVal, ok := finalState.Get("execution_path"); ok { if executionPath, ok := executionPathVal.([]string); ok { execution.ExecutionPath = append(execution.ExecutionPath, executionPath...) @@ -1050,6 +1143,11 @@ Determine if the task is complete or if more actions are needed.`, input, result // Edge condition functions +// shouldAct routes out of the reason node. It is registered on both the +// "act" and "finalize" edges and must name one of them on every path: a +// reasoning step that requests no tool is the ordinary ReAct outcome, and +// returning "" for it left the run with no matching edge and failed with +// ErrNoRoute instead of returning the answer. func (a *BaseAgent) shouldAct(ctx context.Context, state *core.BaseState) (string, error) { // Check if reasoning produced tool calls _, hasToolCalls := state.Get("pending_tool_calls") @@ -1273,6 +1371,9 @@ func (a *BaseAgent) SeedConversation(messages []llm.Message) { for _, m := range messages { a.conversation.AddMessage(m) } + a.mu.Lock() + a.resumeSeeded = len(messages) > 0 + a.mu.Unlock() } // SeedResumeState restores conversation, iteration, and pending tool calls. @@ -1285,6 +1386,7 @@ func (a *BaseAgent) SeedResumeState(messages []llm.Message, iteration int, pendi } a.currentIteration = iteration a.pendingToolCalls = append([]llm.ToolCall(nil), pending...) + a.resumeSeeded = true } // SetGraph sets the agent's execution graph diff --git a/pkg/agent/agent_definition.go b/pkg/agent/agent_definition.go index 8ebae76..5ebafb0 100644 --- a/pkg/agent/agent_definition.go +++ b/pkg/agent/agent_definition.go @@ -468,7 +468,7 @@ func (aad *AdvancedAgentDefinition) CreateAgent() (Agent, error) { // Register custom tools for _, tool := range aad.GetCustomTools() { - aad.toolRegistry.RegisterTool(tool) + _ = aad.toolRegistry.RegisterTool(tool) } // Apply custom graph if available diff --git a/pkg/agent/multi_agent_config.go b/pkg/agent/multi_agent_config.go index 14857fb..4f3b571 100644 --- a/pkg/agent/multi_agent_config.go +++ b/pkg/agent/multi_agent_config.go @@ -8,6 +8,8 @@ package agent import ( "fmt" + "regexp" + "sort" "strings" "time" ) @@ -32,10 +34,47 @@ type RoutingConfig struct { Middleware []MiddlewareConfig `json:"middleware" yaml:"middleware"` } +// Routing match modes. A rule's Match field selects how Pattern is compared +// against the request path. +const ( + MatchExact = "exact" + MatchPrefix = "prefix" + MatchSuffix = "suffix" + MatchContains = "contains" + MatchRegex = "regex" +) + +// Routing condition operators. +const ( + OperatorEquals = "equals" + OperatorContains = "contains" + OperatorPrefix = "prefix" + OperatorSuffix = "suffix" + OperatorRegex = "regex" +) + +// Routing condition sources. +const ( + ConditionHeader = "header" + ConditionQuery = "query" + ConditionIP = "ip" + ConditionMethod = "method" +) + // RoutingRule defines a routing rule type RoutingRule struct { - ID string `json:"id" yaml:"id"` - Pattern string `json:"pattern" yaml:"pattern"` + ID string `json:"id" yaml:"id"` + Pattern string `json:"pattern" yaml:"pattern"` + // Match selects how Pattern is compared against the request path: one of + // exact, prefix, suffix, contains or regex. Empty means exact, which is + // what the HTTP router has always done for path routing. + // + // This field exists because matchesRule used to switch on Pattern itself + // against the literal strings "prefix"/"suffix"/"exact"/"contains", so the + // match mode could never be selected and every real pattern silently fell + // through to prefix matching - disagreeing with the HTTP router, which + // matched the same rule exactly. + Match string `json:"match,omitempty" yaml:"match,omitempty"` AgentID string `json:"agent_id" yaml:"agent_id"` Method string `json:"method" yaml:"method"` Priority int `json:"priority" yaml:"priority"` @@ -43,14 +82,82 @@ type RoutingRule struct { Metadata map[string]interface{} `json:"metadata" yaml:"metadata"` } +// MatchMode returns the normalized match mode for the rule. +func (r RoutingRule) MatchMode() string { + mode := strings.ToLower(strings.TrimSpace(r.Match)) + if mode == "" { + return MatchExact + } + return mode +} + +// MatchesPath reports whether path satisfies the rule's pattern under its +// match mode. An unknown mode never matches; Validate rejects those up front so +// a live system cannot reach this branch with a config it accepted. +func (r RoutingRule) MatchesPath(path string) bool { + switch r.MatchMode() { + case MatchExact: + return path == r.Pattern + case MatchPrefix: + return strings.HasPrefix(path, r.Pattern) + case MatchSuffix: + return strings.HasSuffix(path, r.Pattern) + case MatchContains: + return strings.Contains(path, r.Pattern) + case MatchRegex: + re, err := regexp.Compile(r.Pattern) + if err != nil { + return false + } + return re.MatchString(path) + default: + return false + } +} + // RoutingCondition defines conditions for routing type RoutingCondition struct { - Type string `json:"type" yaml:"type"` // "header", "query", "body", "ip" + Type string `json:"type" yaml:"type"` // "header", "query", "ip", "method" Key string `json:"key" yaml:"key"` Value string `json:"value" yaml:"value"` Operator string `json:"operator" yaml:"operator"` // "equals", "contains", "regex", "prefix", "suffix" } +// OperatorMode returns the normalized operator, defaulting to "equals". +func (rc RoutingCondition) OperatorMode() string { + op := strings.ToLower(strings.TrimSpace(rc.Operator)) + if op == "" { + return OperatorEquals + } + return op +} + +// Evaluate reports whether an actual request value satisfies the condition. +// +// Conditions used to be pure decoration: they were parsed from config and +// carried around on every rule, but nothing ever compared them against a +// request, so a rule guarded by a condition matched every request. +func (rc RoutingCondition) Evaluate(actual string) bool { + switch rc.OperatorMode() { + case OperatorEquals: + return actual == rc.Value + case OperatorContains: + return strings.Contains(actual, rc.Value) + case OperatorPrefix: + return strings.HasPrefix(actual, rc.Value) + case OperatorSuffix: + return strings.HasSuffix(actual, rc.Value) + case OperatorRegex: + re, err := regexp.Compile(rc.Value) + if err != nil { + return false + } + return re.MatchString(actual) + default: + return false + } +} + // MiddlewareConfig defines middleware configuration type MiddlewareConfig struct { Type string `json:"type" yaml:"type"` @@ -584,6 +691,10 @@ func DefaultMultiAgentConfig() *MultiAgentConfig { // Validate validates the multi-agent configuration func (mac *MultiAgentConfig) Validate() error { + if mac == nil { + return fmt.Errorf("multi-agent config is required") + } + if mac.Name == "" { return fmt.Errorf("multi-agent config name is required") } @@ -592,8 +703,18 @@ func (mac *MultiAgentConfig) Validate() error { return fmt.Errorf("at least one agent must be defined") } - // Validate agent configs - for agentID, agentConfig := range mac.Agents { + // Validate agent configs. Iterate in a deterministic order so a config with + // several problems always reports the same one first - map order made the + // reported error vary run to run. + for _, agentID := range sortedKeys(mac.Agents) { + agentConfig := mac.Agents[agentID] + // A YAML key with no body ("agents:\n ghost:") decodes to a nil + // *AgentConfig. Validate used to dereference it immediately and take + // the whole process down with a nil pointer panic instead of reporting + // a config error. + if agentConfig == nil { + return fmt.Errorf("agent %s: configuration is empty", agentID) + } if agentConfig.Name == "" { return fmt.Errorf("agent %s: name is required", agentID) } @@ -606,6 +727,9 @@ func (mac *MultiAgentConfig) Validate() error { if agentConfig.Provider == "" { return fmt.Errorf("agent %s: provider is required", agentID) } + if agentConfig.Timeout < 0 { + return fmt.Errorf("agent %s: timeout must not be negative", agentID) + } } // Validate routing configuration @@ -625,14 +749,42 @@ func (mac *MultiAgentConfig) Validate() error { return nil } +// sortedKeys returns the keys of an agent map in a stable order. +func sortedKeys(m map[string]*AgentConfig) []string { + keys := make([]string, 0, len(m)) + for k := range m { + keys = append(keys, k) + } + sort.Strings(keys) + return keys +} + +// knownRoutingTypes lists the routing types setupRoutingRule can actually +// install. Anything else silently fell through to path routing. +var knownRoutingTypes = map[string]bool{ + "": true, // defaults to path + "path": true, + "host": true, + "header": true, + "query": true, +} + // validateRouting validates routing configuration func (mac *MultiAgentConfig) validateRouting() error { + if !knownRoutingTypes[strings.ToLower(strings.TrimSpace(mac.Routing.Type))] { + return fmt.Errorf("unsupported routing type %q", mac.Routing.Type) + } + if mac.Routing.DefaultAgent != "" { if _, exists := mac.Agents[mac.Routing.DefaultAgent]; !exists { return fmt.Errorf("default agent %s does not exist", mac.Routing.DefaultAgent) } + if !mac.IsAgentEnabled(mac.Routing.DefaultAgent) { + return fmt.Errorf("default agent %s is disabled", mac.Routing.DefaultAgent) + } } + seenIDs := make(map[string]bool, len(mac.Routing.Rules)) for _, rule := range mac.Routing.Rules { if rule.AgentID == "" { return fmt.Errorf("routing rule %s: agent_id is required", rule.ID) @@ -640,11 +792,75 @@ func (mac *MultiAgentConfig) validateRouting() error { if _, exists := mac.Agents[rule.AgentID]; !exists { return fmt.Errorf("routing rule %s: agent %s does not exist", rule.ID, rule.AgentID) } + if !mac.IsAgentEnabled(rule.AgentID) { + return fmt.Errorf("routing rule %s: agent %s is disabled", rule.ID, rule.AgentID) + } + // A rule with no pattern used to be accepted and then installed as a + // route matching nothing (or, for header/query routing, dropped + // without a word), so the agent it names was simply unreachable. + if strings.TrimSpace(rule.Pattern) == "" { + return fmt.Errorf("routing rule %s: pattern is required", rule.ID) + } + if rule.ID != "" { + if seenIDs[rule.ID] { + return fmt.Errorf("routing rule %s: duplicate rule id", rule.ID) + } + seenIDs[rule.ID] = true + } + if err := validateMatchMode(rule); err != nil { + return fmt.Errorf("routing rule %s: %w", rule.ID, err) + } + for i, cond := range rule.Conditions { + if err := validateCondition(cond); err != nil { + return fmt.Errorf("routing rule %s: condition %d: %w", rule.ID, i, err) + } + } } return nil } +func validateMatchMode(rule RoutingRule) error { + switch rule.MatchMode() { + case MatchExact, MatchPrefix, MatchSuffix, MatchContains: + return nil + case MatchRegex: + if _, err := regexp.Compile(rule.Pattern); err != nil { + return fmt.Errorf("invalid regex pattern: %w", err) + } + return nil + default: + return fmt.Errorf("unsupported match mode %q", rule.Match) + } +} + +func validateCondition(cond RoutingCondition) error { + switch strings.ToLower(strings.TrimSpace(cond.Type)) { + case ConditionHeader, ConditionQuery: + if cond.Key == "" { + return fmt.Errorf("%s condition requires a key", cond.Type) + } + case ConditionIP, ConditionMethod: + // Keyless: the value is taken from the request itself. + default: + // Anything else (notably "body") cannot be evaluated by the router, so + // accepting it would mean silently ignoring a security-relevant guard. + return fmt.Errorf("unsupported condition type %q", cond.Type) + } + + switch cond.OperatorMode() { + case OperatorEquals, OperatorContains, OperatorPrefix, OperatorSuffix: + return nil + case OperatorRegex: + if _, err := regexp.Compile(cond.Value); err != nil { + return fmt.Errorf("invalid regex value: %w", err) + } + return nil + default: + return fmt.Errorf("unsupported operator %q", cond.Operator) + } +} + // validateDeployment validates deployment configuration func (mac *MultiAgentConfig) validateDeployment() error { if mac.Deployment.Type == "" { @@ -666,17 +882,65 @@ func (mac *MultiAgentConfig) validateDeployment() error { } } + if hc := mac.Deployment.HealthCheck; hc != nil && hc.Enabled { + if err := hc.validate("health_check"); err != nil { + return err + } + for agentID, specific := range hc.AgentSpecific { + if specific == nil { + return fmt.Errorf("health_check.agent_specific.%s: configuration is empty", agentID) + } + if _, exists := mac.Agents[agentID]; !exists { + return fmt.Errorf("health_check.agent_specific.%s: agent does not exist", agentID) + } + if err := specific.validate("health_check.agent_specific." + agentID); err != nil { + return err + } + } + } + + if sc := mac.Deployment.Scaling; sc != nil && sc.Enabled { + if sc.MinReplicas < 1 { + return fmt.Errorf("scaling min_replicas must be at least 1") + } + if sc.MaxReplicas < sc.MinReplicas { + return fmt.Errorf("scaling max_replicas (%d) must not be below min_replicas (%d)", sc.MaxReplicas, sc.MinReplicas) + } + } + return nil } -// GetAgentByPath returns the agent ID for a given path +// SortedRules returns the routing rules ordered the way the HTTP router +// installs them: highest priority first, ties broken by rule ID so the order is +// stable across runs. +func (mac *MultiAgentConfig) SortedRules() []RoutingRule { + if mac.Routing == nil { + return nil + } + rules := make([]RoutingRule, len(mac.Routing.Rules)) + copy(rules, mac.Routing.Rules) + sort.SliceStable(rules, func(i, j int) bool { + if rules[i].Priority != rules[j].Priority { + return rules[i].Priority > rules[j].Priority + } + return rules[i].ID < rules[j].ID + }) + return rules +} + +// GetAgentByPath returns the agent ID for a given path. +// +// Rules are consulted highest-priority-first. They used to be walked in raw +// declaration order, so this answered with a different agent than the HTTP +// router - which does sort by priority - whenever two patterns overlapped. func (mac *MultiAgentConfig) GetAgentByPath(path string) (string, bool) { if mac.Routing == nil { return "", false } // Check routing rules - for _, rule := range mac.Routing.Rules { + for _, rule := range mac.SortedRules() { if mac.matchesRule(rule, path) { return rule.AgentID, true } @@ -692,19 +956,7 @@ func (mac *MultiAgentConfig) GetAgentByPath(path string) (string, bool) { // matchesRule checks if a path matches a routing rule func (mac *MultiAgentConfig) matchesRule(rule RoutingRule, path string) bool { - switch rule.Pattern { - case "prefix": - return strings.HasPrefix(path, rule.Pattern) - case "suffix": - return strings.HasSuffix(path, rule.Pattern) - case "exact": - return path == rule.Pattern - case "contains": - return strings.Contains(path, rule.Pattern) - default: - // Default to prefix matching - return strings.HasPrefix(path, rule.Pattern) - } + return rule.MatchesPath(path) } // GetAgentPaths returns all paths for each agent @@ -746,13 +998,257 @@ func (mac *MultiAgentConfig) ListAgentIDs() []string { return ids } +// IsAgentEnabled reports whether an agent should be started and routed to. +// +// An agent is disabled by setting metadata.disabled: true (or +// metadata.enabled: false) on its config. GetEnabledAgents previously returned +// every agent with a comment admitting it did no filtering at all, so a config +// that switched an agent off still got it created, routed to and health +// checked. +func (mac *MultiAgentConfig) IsAgentEnabled(agentID string) bool { + config, exists := mac.Agents[agentID] + if !exists || config == nil { + return false + } + if disabled, ok := boolMetadata(config.Metadata, "disabled"); ok { + return !disabled + } + if enabled, ok := boolMetadata(config.Metadata, "enabled"); ok { + return enabled + } + return true +} + +// boolMetadata reads a boolean from an agent's metadata, tolerating the string +// forms YAML and JSON configs produce ("true"/"false"). +func boolMetadata(metadata map[string]interface{}, key string) (bool, bool) { + raw, exists := metadata[key] + if !exists { + return false, false + } + switch v := raw.(type) { + case bool: + return v, true + case string: + switch strings.ToLower(strings.TrimSpace(v)) { + case "true", "yes", "1": + return true, true + case "false", "no", "0": + return false, true + } + } + return false, false +} + // GetEnabledAgents returns only enabled agents func (mac *MultiAgentConfig) GetEnabledAgents() map[string]*AgentConfig { enabled := make(map[string]*AgentConfig) for id, config := range mac.Agents { - // Check if agent is disabled (assuming we extend AgentConfig with a Disabled field) - // For now, all agents are considered enabled + if !mac.IsAgentEnabled(id) { + continue + } enabled[id] = config } return enabled } + +// Health check defaults. A health check block that enables checking without +// saying how often used to reach time.NewTicker(0), which panics on a +// background goroutine and takes the whole process down. +const ( + DefaultHealthCheckPeriod = 10 * time.Second + DefaultHealthCheckTimeout = 5 * time.Second + DefaultHealthCheckFailureThreshold = 3 +) + +// validate rejects health check settings that cannot be honored. +func (hc *HealthCheckConfig) validate(field string) error { + if hc.PeriodSeconds < 0 { + return fmt.Errorf("%s: period_seconds must not be negative", field) + } + if hc.InitialDelaySeconds < 0 { + return fmt.Errorf("%s: initial_delay_seconds must not be negative", field) + } + if hc.TimeoutSeconds < 0 { + return fmt.Errorf("%s: timeout_seconds must not be negative", field) + } + if hc.FailureThreshold < 0 { + return fmt.Errorf("%s: failure_threshold must not be negative", field) + } + return nil +} + +// Period returns the interval between checks, substituting a safe default for +// an unset value rather than handing 0 to time.NewTicker. +func (hc *HealthCheckConfig) Period() time.Duration { + if hc == nil || hc.PeriodSeconds <= 0 { + return DefaultHealthCheckPeriod + } + return time.Duration(hc.PeriodSeconds) * time.Second +} + +// InitialDelay returns how long to wait before the first check. +func (hc *HealthCheckConfig) InitialDelay() time.Duration { + if hc == nil || hc.InitialDelaySeconds <= 0 { + return 0 + } + return time.Duration(hc.InitialDelaySeconds) * time.Second +} + +// Timeout returns the per-check timeout. +func (hc *HealthCheckConfig) Timeout() time.Duration { + if hc == nil || hc.TimeoutSeconds <= 0 { + return DefaultHealthCheckTimeout + } + return time.Duration(hc.TimeoutSeconds) * time.Second +} + +// Failures returns the consecutive-failure count that marks an agent unhealthy. +func (hc *HealthCheckConfig) Failures() int { + if hc == nil || hc.FailureThreshold <= 0 { + return DefaultHealthCheckFailureThreshold + } + return hc.FailureThreshold +} + +// RedactedPlaceholder replaces every secret value in a redacted config. +const RedactedPlaceholder = "[REDACTED]" + +// Redacted returns a copy of the configuration with every credential replaced +// by RedactedPlaceholder. +// +// The /config endpoint served the raw struct, so a plain unauthenticated GET +// returned provider API keys, database and cache passwords, SMTP credentials, +// Slack webhook URLs and both secret maps verbatim. +func (mac *MultiAgentConfig) Redacted() *MultiAgentConfig { + if mac == nil { + return nil + } + + safe := *mac + + if mac.Deployment != nil { + deployment := *mac.Deployment + deployment.Secrets = redactMap(mac.Deployment.Secrets) + safe.Deployment = &deployment + } + + if mac.Shared != nil { + shared := *mac.Shared + shared.Secrets = redactMap(mac.Shared.Secrets) + + if mac.Shared.LLMProviders != nil { + providers := make(map[string]*LLMProviderConfig, len(mac.Shared.LLMProviders)) + for name, provider := range mac.Shared.LLMProviders { + if provider == nil { + providers[name] = nil + continue + } + copied := *provider + if copied.APIKey != "" { + copied.APIKey = RedactedPlaceholder + } + copied.Config = redactAnyMap(provider.Config) + providers[name] = &copied + } + shared.LLMProviders = providers + } + + if mac.Shared.Database != nil { + db := *mac.Shared.Database + if db.Password != "" { + db.Password = RedactedPlaceholder + } + shared.Database = &db + } + + if mac.Shared.Cache != nil { + cache := *mac.Shared.Cache + if cache.Password != "" { + cache.Password = RedactedPlaceholder + } + shared.Cache = &cache + } + + if mac.Shared.Security != nil { + security := *mac.Shared.Security + if mac.Shared.Security.Authentication != nil { + auth := *mac.Shared.Security.Authentication + auth.Config = redactAnyMap(mac.Shared.Security.Authentication.Config) + security.Authentication = &auth + } + if mac.Shared.Security.Authorization != nil { + authz := *mac.Shared.Security.Authorization + authz.Config = redactAnyMap(mac.Shared.Security.Authorization.Config) + security.Authorization = &authz + } + shared.Security = &security + } + + if mac.Shared.Monitoring != nil && mac.Shared.Monitoring.Alerting != nil { + monitoring := *mac.Shared.Monitoring + alerting := *mac.Shared.Monitoring.Alerting + if alerting.Slack != nil { + slack := *alerting.Slack + if slack.WebhookURL != "" { + slack.WebhookURL = RedactedPlaceholder + } + alerting.Slack = &slack + } + if alerting.Email != nil && alerting.Email.SMTP != nil { + email := *alerting.Email + smtp := *alerting.Email.SMTP + if smtp.Password != "" { + smtp.Password = RedactedPlaceholder + } + email.SMTP = &smtp + alerting.Email = &email + } + monitoring.Alerting = &alerting + shared.Monitoring = &monitoring + } + + safe.Shared = &shared + } + + return &safe +} + +// secretKeyFragments matches config keys whose values must never be echoed. +var secretKeyFragments = []string{"secret", "password", "passwd", "token", "api_key", "apikey", "credential", "private"} + +func looksSecret(key string) bool { + lower := strings.ToLower(key) + for _, fragment := range secretKeyFragments { + if strings.Contains(lower, fragment) { + return true + } + } + return false +} + +func redactMap(in map[string]string) map[string]string { + if in == nil { + return nil + } + out := make(map[string]string, len(in)) + for k := range in { + out[k] = RedactedPlaceholder + } + return out +} + +func redactAnyMap(in map[string]interface{}) map[string]interface{} { + if in == nil { + return nil + } + out := make(map[string]interface{}, len(in)) + for k, v := range in { + if looksSecret(k) { + out[k] = RedactedPlaceholder + continue + } + out[k] = v + } + return out +} diff --git a/pkg/agent/multi_agent_defects_test.go b/pkg/agent/multi_agent_defects_test.go new file mode 100644 index 0000000..90e8ccd --- /dev/null +++ b/pkg/agent/multi_agent_defects_test.go @@ -0,0 +1,994 @@ +// Copyright (c) 2024 GoLangGraph Team +// +// Licensed under the MIT License. See LICENSE file in the project root for full license information. + +package agent + +import ( + "context" + "encoding/json" + "errors" + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + yaml "gopkg.in/yaml.v3" + + "github.com/UnicoLab/GoLangGraph/pkg/llm" + "github.com/UnicoLab/GoLangGraph/pkg/tools" + "github.com/UnicoLab/GoLangGraph/test/fakes" +) + +// --------------------------------------------------------------------------- +// helpers +// --------------------------------------------------------------------------- + +// newTestAgentConfig builds a minimal valid agent config bound to provider. +func newTestAgentConfig(id, provider string) *AgentConfig { + return &AgentConfig{ + ID: id, + Name: id, + Type: AgentTypeChat, + Model: "fake-model", + Provider: provider, + SystemPrompt: "you are a test agent", + MaxIterations: 1, + Tools: []string{}, + } +} + +// newManagerWithProviders builds a manager whose LLM manager already knows the +// supplied providers. It stops the manager when the test finishes so no test +// can leave health checkers running behind it. +func newManagerWithProviders(t *testing.T, config *MultiAgentConfig, providers map[string]*fakes.Provider) *MultiAgentManager { + t.Helper() + + llmManager := llm.NewProviderManager() + for name, provider := range providers { + require.NoError(t, llmManager.RegisterProvider(name, provider)) + } + + manager, err := NewMultiAgentManager(config, llmManager, tools.NewToolRegistry()) + require.NoError(t, err) + t.Cleanup(func() { + stopCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + _ = manager.Stop(stopCtx) + }) + return manager +} + +// newSingleAgentManager is the common one-agent, one-route setup. +func newSingleAgentManager(t *testing.T, mutate func(*MultiAgentConfig)) (*MultiAgentManager, *fakes.Provider) { + t.Helper() + + provider := fakes.NewProvider("fake", "hello from the fake provider") + config := &MultiAgentConfig{ + Name: "single-agent", + Version: "1.0", + Agents: map[string]*AgentConfig{"solo": newTestAgentConfig("solo", "fake")}, + Routing: &RoutingConfig{ + Type: "path", + Rules: []RoutingRule{{ID: "solo-rule", Pattern: "/solo", AgentID: "solo", Method: "POST"}}, + }, + } + if mutate != nil { + mutate(config) + } + + return newManagerWithProviders(t, config, map[string]*fakes.Provider{"fake": provider}), provider +} + +func postInput(t *testing.T, url, input string) (int, string) { + t.Helper() + body := `{"input":` + jsonString(input) + `}` + resp, err := http.Post(url, "application/json", strings.NewReader(body)) + require.NoError(t, err) + defer func() { _ = resp.Body.Close() }() + payload, err := io.ReadAll(resp.Body) + require.NoError(t, err) + return resp.StatusCode, string(payload) +} + +// jsonString quotes a value for embedding in a JSON request body. +func jsonString(s string) string { + encoded, err := json.Marshal(s) + if err != nil { + return `""` + } + return string(encoded) +} + +func getJSON(t *testing.T, url string) (int, map[string]interface{}) { + t.Helper() + resp, err := http.Get(url) + require.NoError(t, err) + defer func() { _ = resp.Body.Close() }() + + var decoded map[string]interface{} + payload, err := io.ReadAll(resp.Body) + require.NoError(t, err) + if len(payload) > 0 { + _ = json.Unmarshal(payload, &decoded) + } + return resp.StatusCode, decoded +} + +// --------------------------------------------------------------------------- +// Config validation defects +// --------------------------------------------------------------------------- + +// Defect: MultiAgentConfig.Validate dereferenced every *AgentConfig without a +// nil check, so a YAML agent key with no body ("agents:\n ghost:") took the +// process down with a nil pointer panic instead of returning a config error. +func TestRegression_ValidateRejectsEmptyAgentInsteadOfPanicking(t *testing.T) { + var config MultiAgentConfig + require.NoError(t, yaml.Unmarshal([]byte("name: ghosts\nagents:\n ghost:\n"), &config)) + require.Contains(t, config.Agents, "ghost") + require.Nil(t, config.Agents["ghost"], "the YAML must actually decode to a nil agent config") + + err := config.Validate() + + require.Error(t, err) + assert.Contains(t, err.Error(), "agent ghost: configuration is empty") +} + +// Defect: the manager panicked in setupRouting on a config with no routing +// block, even though Validate explicitly treats routing as optional. +func TestRegression_ManagerBuildsWithoutRoutingSection(t *testing.T) { + config := &MultiAgentConfig{ + Name: "no-routing", + Agents: map[string]*AgentConfig{"solo": newTestAgentConfig("solo", "fake")}, + // Routing deliberately nil. + } + + manager := newManagerWithProviders(t, config, map[string]*fakes.Provider{ + "fake": fakes.NewProvider("fake", "hi"), + }) + + // Management endpoints must still be reachable. + server := httptest.NewServer(manager.GetRouter()) + defer server.Close() + + status, body := getJSON(t, server.URL+"/routing") + assert.Equal(t, http.StatusOK, status) + assert.NotNil(t, body) +} + +// Defect: validateRouting accepted rules with no pattern, which then installed +// a route that could never match, leaving the named agent unreachable. +func TestRegression_ValidateRejectsPatternlessRule(t *testing.T) { + config := &MultiAgentConfig{ + Name: "patternless", + Agents: map[string]*AgentConfig{"solo": newTestAgentConfig("solo", "fake")}, + Routing: &RoutingConfig{ + Type: "path", + Rules: []RoutingRule{{ID: "broken", AgentID: "solo"}}, + }, + } + + err := config.Validate() + + require.Error(t, err) + assert.Contains(t, err.Error(), "pattern is required") +} + +// Defect: Validate accepted any routing type and setupRoutingRule silently fell +// through to path routing for the unknown ones. +func TestRegression_ValidateRejectsUnknownRoutingType(t *testing.T) { + config := &MultiAgentConfig{ + Name: "weird-routing", + Agents: map[string]*AgentConfig{"solo": newTestAgentConfig("solo", "fake")}, + Routing: &RoutingConfig{ + Type: "carrier-pigeon", + Rules: []RoutingRule{{ID: "r", Pattern: "/solo", AgentID: "solo"}}, + }, + } + + err := config.Validate() + + require.Error(t, err) + assert.Contains(t, err.Error(), "unsupported routing type") +} + +// --------------------------------------------------------------------------- +// Routing defects +// --------------------------------------------------------------------------- + +// Defect: matchesRule switched on rule.Pattern against the literal strings +// "prefix"/"suffix"/"exact"/"contains", so the match mode could never be +// chosen and every real pattern fell through to prefix matching - while the +// HTTP router matched the same rule exactly. GetAgentByPath also ignored +// Priority, so it disagreed with the router about which rule wins. +func TestRegression_GetAgentByPathHonorsPriorityAndMatchMode(t *testing.T) { + config := &MultiAgentConfig{ + Name: "priority", + Agents: map[string]*AgentConfig{ + "broad": newTestAgentConfig("broad", "fake"), + "specific": newTestAgentConfig("specific", "fake"), + "suffixing": newTestAgentConfig("suffixing", "fake"), + }, + Routing: &RoutingConfig{ + Type: "path", + Rules: []RoutingRule{ + {ID: "broad", Pattern: "/api", Match: MatchPrefix, AgentID: "broad", Priority: 1}, + {ID: "specific", Pattern: "/api/special", Match: MatchPrefix, AgentID: "specific", Priority: 100}, + {ID: "suffixing", Pattern: ".json", Match: MatchSuffix, AgentID: "suffixing", Priority: 50}, + }, + }, + } + require.NoError(t, config.Validate()) + + // Highest priority wins even though it is declared second. + agentID, found := config.GetAgentByPath("/api/special/thing") + assert.True(t, found) + assert.Equal(t, "specific", agentID, "priority must decide, not declaration order") + + // The lower-priority prefix rule still serves everything else under /api. + agentID, found = config.GetAgentByPath("/api/other") + assert.True(t, found) + assert.Equal(t, "broad", agentID) + + // Suffix mode is now reachable at all. + agentID, found = config.GetAgentByPath("/reports/latest.json") + assert.True(t, found) + assert.Equal(t, "suffixing", agentID) + + // A pattern that happens to be spelled like a match mode is treated as a + // pattern, which is what the old switch got wrong. + odd := &MultiAgentConfig{ + Name: "odd", + Agents: map[string]*AgentConfig{"solo": newTestAgentConfig("solo", "fake")}, + Routing: &RoutingConfig{ + Type: "path", + Rules: []RoutingRule{{ID: "r", Pattern: "exact", Match: MatchExact, AgentID: "solo"}}, + }, + } + _, found = odd.GetAgentByPath("nothing-like-it") + assert.False(t, found, "an unrelated path must not match the literal pattern \"exact\"") + agentID, found = odd.GetAgentByPath("exact") + assert.True(t, found) + assert.Equal(t, "solo", agentID) +} + +// Defect: header patterns were split on every colon and dropped unless exactly +// two parts came back, so "Authorization: Bearer xyz" - a value containing a +// colon-separated scheme - installed no route at all, silently, and the agent +// was unreachable while startup reported success. +func TestRegression_HeaderRoutingRuleKeepsColonsInValue(t *testing.T) { + provider := fakes.NewProvider("fake", "routed") + config := &MultiAgentConfig{ + Name: "header-routing", + Agents: map[string]*AgentConfig{"solo": newTestAgentConfig("solo", "fake")}, + Routing: &RoutingConfig{ + Type: "header", + Rules: []RoutingRule{{ID: "bearer", Pattern: "Authorization: Bearer xyz", AgentID: "solo", Method: "POST"}}, + }, + } + + manager := newManagerWithProviders(t, config, map[string]*fakes.Provider{"fake": provider}) + server := httptest.NewServer(manager.GetRouter()) + defer server.Close() + + req, err := http.NewRequest(http.MethodPost, server.URL+"/anything", strings.NewReader(`{"input":"hi"}`)) + require.NoError(t, err) + req.Header.Set("Authorization", "Bearer xyz") + + resp, err := http.DefaultClient.Do(req) + require.NoError(t, err) + defer func() { _ = resp.Body.Close() }() + + assert.Equal(t, http.StatusOK, resp.StatusCode, "the rule must actually be installed") + assert.Equal(t, 1, provider.Calls()) + + // A request without the header must not reach the agent. + resp2, err := http.Post(server.URL+"/anything", "application/json", strings.NewReader(`{"input":"hi"}`)) + require.NoError(t, err) + defer func() { _ = resp2.Body.Close() }() + assert.Equal(t, http.StatusNotFound, resp2.StatusCode) +} + +// Defect: a header/query pattern the router could not express left `route` nil +// and setupRoutingRule returned without a word, so a broken config started +// cleanly with a silently missing route. Failures are now reported. +func TestRegression_UninstallableRoutingRuleIsReported(t *testing.T) { + config := &MultiAgentConfig{ + Name: "bad-header-pattern", + Agents: map[string]*AgentConfig{"solo": newTestAgentConfig("solo", "fake")}, + Routing: &RoutingConfig{ + Type: "header", + Rules: []RoutingRule{{ID: "no-colon", Pattern: "AuthorizationBearerXyz", AgentID: "solo"}}, + }, + } + + llmManager := llm.NewProviderManager() + require.NoError(t, llmManager.RegisterProvider("fake", fakes.NewProvider("fake", "hi"))) + + _, err := NewMultiAgentManager(config, llmManager, tools.NewToolRegistry()) + + require.Error(t, err) + assert.Contains(t, err.Error(), "no-colon") +} + +// Defect: RoutingRule.Conditions were parsed, stored and serialized but never +// compared against a request, so a rule guarded by a condition matched +// everything. +func TestRegression_RoutingConditionsAreEvaluated(t *testing.T) { + provider := fakes.NewProvider("fake", "conditional") + config := &MultiAgentConfig{ + Name: "conditions", + Agents: map[string]*AgentConfig{"solo": newTestAgentConfig("solo", "fake")}, + Routing: &RoutingConfig{ + Type: "path", + Rules: []RoutingRule{{ + ID: "guarded", + Pattern: "/guarded", + AgentID: "solo", + Method: "POST", + Conditions: []RoutingCondition{ + {Type: ConditionHeader, Key: "X-Tenant", Value: "acme", Operator: OperatorEquals}, + }, + }}, + }, + } + + manager := newManagerWithProviders(t, config, map[string]*fakes.Provider{"fake": provider}) + server := httptest.NewServer(manager.GetRouter()) + defer server.Close() + + // Condition satisfied. + req, err := http.NewRequest(http.MethodPost, server.URL+"/guarded", strings.NewReader(`{"input":"hi"}`)) + require.NoError(t, err) + req.Header.Set("X-Tenant", "acme") + resp, err := http.DefaultClient.Do(req) + require.NoError(t, err) + defer func() { _ = resp.Body.Close() }() + assert.Equal(t, http.StatusOK, resp.StatusCode) + + // Condition violated: the rule must not fire. + req2, err := http.NewRequest(http.MethodPost, server.URL+"/guarded", strings.NewReader(`{"input":"hi"}`)) + require.NoError(t, err) + req2.Header.Set("X-Tenant", "someone-else") + resp2, err := http.DefaultClient.Do(req2) + require.NoError(t, err) + defer func() { _ = resp2.Body.Close() }() + assert.Equal(t, http.StatusNotFound, resp2.StatusCode) + + assert.Equal(t, 1, provider.Calls(), "only the satisfying request may reach the agent") +} + +// Defect: a "body" condition can never be evaluated by a router, and accepting +// it meant silently ignoring a guard the operator wrote down. +func TestRegression_UnsupportedConditionTypeIsRejected(t *testing.T) { + config := &MultiAgentConfig{ + Name: "body-condition", + Agents: map[string]*AgentConfig{"solo": newTestAgentConfig("solo", "fake")}, + Routing: &RoutingConfig{ + Type: "path", + Rules: []RoutingRule{{ + ID: "r", + Pattern: "/solo", + AgentID: "solo", + Conditions: []RoutingCondition{{Type: "body", Key: "kind", Value: "x"}}, + }}, + }, + } + + err := config.Validate() + + require.Error(t, err) + assert.Contains(t, err.Error(), "unsupported condition type") +} + +// --------------------------------------------------------------------------- +// Security defects +// --------------------------------------------------------------------------- + +// Defect: GET /config encoded the raw configuration, so an unauthenticated +// caller got provider API keys, database and cache passwords and both secret +// maps back verbatim. +func TestRegression_ConfigEndpointRedactsSecrets(t *testing.T) { + manager, _ := newSingleAgentManager(t, func(config *MultiAgentConfig) { + config.Deployment = &DeploymentConfig{ + Type: "docker", + Replicas: 1, + Secrets: map[string]string{"tls": "DEPLOYMENT-SECRET"}, + } + config.Shared = &SharedConfig{ + Secrets: map[string]string{"token": "SHARED-SECRET"}, + LLMProviders: map[string]*LLMProviderConfig{"openai": {Type: "openai", APIKey: "sk-LEAKED-KEY"}}, + Database: &DatabaseConfig{Host: "db", Password: "DB-PASSWORD"}, + Cache: &CacheConfig{Host: "redis", Password: "CACHE-PASSWORD"}, + Security: &SecurityConfig{ + Authentication: &AuthConfig{Type: "apikey", Config: map[string]interface{}{"api_keys": []interface{}{"AUTH-KEY"}}}, + }, + Monitoring: &MonitoringConfig{Alerting: &AlertingConfig{ + Slack: &SlackConfig{WebhookURL: "https://hooks.example/SLACK-SECRET"}, + Email: &EmailConfig{SMTP: &SMTPConfig{Host: "smtp", Password: "SMTP-PASSWORD"}}, + }}, + } + }) + + server := httptest.NewServer(manager.GetRouter()) + defer server.Close() + + resp, err := http.Get(server.URL + "/config") + require.NoError(t, err) + defer func() { _ = resp.Body.Close() }() + payload, err := io.ReadAll(resp.Body) + require.NoError(t, err) + body := string(payload) + + for _, secret := range []string{ + "DEPLOYMENT-SECRET", "SHARED-SECRET", "sk-LEAKED-KEY", "DB-PASSWORD", + "CACHE-PASSWORD", "AUTH-KEY", "SLACK-SECRET", "SMTP-PASSWORD", + } { + assert.NotContains(t, body, secret, "/config must not echo %s", secret) + } + assert.Contains(t, body, RedactedPlaceholder) + + // Redaction must not mutate the manager's live configuration. + assert.Equal(t, "sk-LEAKED-KEY", manager.GetConfig().Shared.LLMProviders["openai"].APIKey) + assert.Equal(t, "DB-PASSWORD", manager.GetConfig().Shared.Database.Password) +} + +// Defect: the auth middleware read an API key and then accepted any non-empty +// value ("in a real implementation, validate the API key"), so a deployment +// that believed it required authentication was open to anyone who sent a +// header at all. +func TestRegression_AuthMiddlewareValidatesTheKey(t *testing.T) { + provider := fakes.NewProvider("fake", "authorized") + config := &MultiAgentConfig{ + Name: "auth", + Agents: map[string]*AgentConfig{"solo": newTestAgentConfig("solo", "fake")}, + Routing: &RoutingConfig{ + Type: "path", + Rules: []RoutingRule{{ID: "r", Pattern: "/solo", AgentID: "solo", Method: "POST"}}, + Middleware: []MiddlewareConfig{{ + Type: "auth", + Enabled: true, + Config: map[string]interface{}{"api_keys": []interface{}{"correct-horse"}}, + }}, + }, + } + + manager := newManagerWithProviders(t, config, map[string]*fakes.Provider{"fake": provider}) + server := httptest.NewServer(manager.GetRouter()) + defer server.Close() + + cases := []struct { + name string + key string + expected int + }{ + {"no key", "", http.StatusUnauthorized}, + {"wrong key", "battery-staple", http.StatusUnauthorized}, + {"correct key", "correct-horse", http.StatusOK}, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + req, err := http.NewRequest(http.MethodPost, server.URL+"/solo", strings.NewReader(`{"input":"hi"}`)) + require.NoError(t, err) + if tc.key != "" { + req.Header.Set("X-API-Key", tc.key) + } + resp, err := http.DefaultClient.Do(req) + require.NoError(t, err) + defer func() { _ = resp.Body.Close() }() + assert.Equal(t, tc.expected, resp.StatusCode) + }) + } + + assert.Equal(t, 1, provider.Calls(), "only the authorized request may reach the agent") +} + +// Defect: enabling the auth middleware with no keys configured used to accept +// every request. Failing construction is the honest outcome - a config that +// asks for authentication it cannot perform must not start. +func TestRegression_AuthMiddlewareWithoutKeysFailsToStart(t *testing.T) { + config := &MultiAgentConfig{ + Name: "auth-no-keys", + Agents: map[string]*AgentConfig{"solo": newTestAgentConfig("solo", "fake")}, + Routing: &RoutingConfig{ + Type: "path", + Rules: []RoutingRule{{ID: "r", Pattern: "/solo", AgentID: "solo", Method: "POST"}}, + Middleware: []MiddlewareConfig{{Type: "auth", Enabled: true}}, + }, + } + + llmManager := llm.NewProviderManager() + require.NoError(t, llmManager.RegisterProvider("fake", fakes.NewProvider("fake", "hi"))) + + _, err := NewMultiAgentManager(config, llmManager, tools.NewToolRegistry()) + + require.Error(t, err) + assert.Contains(t, err.Error(), "no API keys are configured") +} + +// Defect: rateLimitMiddleware called next.ServeHTTP and nothing else while the +// config carried a complete RateLimitConfig, so every configured limit was +// silently ignored. +func TestRegression_RateLimitMiddlewareEnforcesGlobalLimit(t *testing.T) { + provider := fakes.NewProvider("fake", "limited") + config := &MultiAgentConfig{ + Name: "rate-limited", + Agents: map[string]*AgentConfig{"solo": newTestAgentConfig("solo", "fake")}, + Routing: &RoutingConfig{ + Type: "path", + Rules: []RoutingRule{{ID: "r", Pattern: "/solo", AgentID: "solo", Method: "POST"}}, + Middleware: []MiddlewareConfig{{Type: "rate_limit", Enabled: true}}, + }, + Shared: &SharedConfig{Security: &SecurityConfig{ + RateLimit: &RateLimitConfig{ + Enabled: true, + Global: &RateLimit{Requests: 3, Period: time.Hour, Burst: 3}, + }, + }}, + } + + manager := newManagerWithProviders(t, config, map[string]*fakes.Provider{"fake": provider}) + server := httptest.NewServer(manager.GetRouter()) + defer server.Close() + + allowed, limited := 0, 0 + for i := 0; i < 8; i++ { + status, _ := postInput(t, server.URL+"/solo", "hi") + switch status { + case http.StatusOK: + allowed++ + case http.StatusTooManyRequests: + limited++ + default: + t.Fatalf("unexpected status %d", status) + } + } + + assert.Equal(t, 3, allowed, "the configured budget of 3 requests must be honored") + assert.Equal(t, 5, limited) + assert.Equal(t, 3, provider.Calls(), "rejected requests must not reach the agent") +} + +// The per-agent budget in RateLimitConfig.PerAgent needs the agent's identity, +// which only the agent handler has; it was previously never consulted at all. +func TestRegression_RateLimitAppliesPerAgentBudget(t *testing.T) { + busy := fakes.NewProvider("busy", "busy reply") + calm := fakes.NewProvider("calm", "calm reply") + + config := &MultiAgentConfig{ + Name: "per-agent-rate-limit", + Agents: map[string]*AgentConfig{ + "busy": newTestAgentConfig("busy", "busy"), + "calm": newTestAgentConfig("calm", "calm"), + }, + Routing: &RoutingConfig{ + Type: "path", + Rules: []RoutingRule{ + {ID: "busy", Pattern: "/busy", AgentID: "busy", Method: "POST"}, + {ID: "calm", Pattern: "/calm", AgentID: "calm", Method: "POST"}, + }, + Middleware: []MiddlewareConfig{{Type: "rate_limit", Enabled: true}}, + }, + Shared: &SharedConfig{Security: &SecurityConfig{ + RateLimit: &RateLimitConfig{ + Enabled: true, + PerAgent: map[string]*RateLimit{"busy": {Requests: 2, Period: time.Hour, Burst: 2}}, + }, + }}, + } + + manager := newManagerWithProviders(t, config, map[string]*fakes.Provider{"busy": busy, "calm": calm}) + server := httptest.NewServer(manager.GetRouter()) + defer server.Close() + + busyLimited := 0 + for i := 0; i < 5; i++ { + if status, _ := postInput(t, server.URL+"/busy", "hi"); status == http.StatusTooManyRequests { + busyLimited++ + } + } + assert.Equal(t, 3, busyLimited, "the busy agent's budget of 2 must be enforced") + + // The agent without a budget is untouched. + for i := 0; i < 5; i++ { + status, _ := postInput(t, server.URL+"/calm", "hi") + assert.Equal(t, http.StatusOK, status) + } + assert.Equal(t, 5, calm.Calls()) +} + +// Defect: an unknown middleware type was logged as a warning and dropped, so a +// typo in a config silently disabled the protection it named. +func TestRegression_UnknownMiddlewareTypeFailsToStart(t *testing.T) { + config := &MultiAgentConfig{ + Name: "typo-middleware", + Agents: map[string]*AgentConfig{"solo": newTestAgentConfig("solo", "fake")}, + Routing: &RoutingConfig{ + Type: "path", + Rules: []RoutingRule{{ID: "r", Pattern: "/solo", AgentID: "solo", Method: "POST"}}, + Middleware: []MiddlewareConfig{{Type: "authh", Enabled: true}}, + }, + } + + llmManager := llm.NewProviderManager() + require.NoError(t, llmManager.RegisterProvider("fake", fakes.NewProvider("fake", "hi"))) + + _, err := NewMultiAgentManager(config, llmManager, tools.NewToolRegistry()) + + require.Error(t, err) + assert.Contains(t, err.Error(), "unknown middleware type") +} + +// Defect: corsMiddleware read mam.config.Shared.Security.CORS.Enabled with no +// nil check on any level of that chain, so a config that enabled the CORS +// middleware without a shared security section panicked inside the HTTP +// handler - after the server was already accepting traffic. +func TestRegression_CORSMiddlewareSurvivesMissingSharedSection(t *testing.T) { + provider := fakes.NewProvider("fake", "cors ok") + config := &MultiAgentConfig{ + Name: "cors-no-shared", + Agents: map[string]*AgentConfig{"solo": newTestAgentConfig("solo", "fake")}, + Routing: &RoutingConfig{ + Type: "path", + Rules: []RoutingRule{{ID: "r", Pattern: "/solo", AgentID: "solo", Method: "POST"}}, + Middleware: []MiddlewareConfig{{Type: "cors", Enabled: true}}, + }, + // Shared deliberately nil. + } + + manager := newManagerWithProviders(t, config, map[string]*fakes.Provider{"fake": provider}) + server := httptest.NewServer(manager.GetRouter()) + defer server.Close() + + status, _ := postInput(t, server.URL+"/solo", "hi") + assert.Equal(t, http.StatusOK, status) +} + +// With CORS actually configured the headers must be emitted. +func TestCORSMiddlewareEmitsConfiguredHeaders(t *testing.T) { + manager, _ := newSingleAgentManager(t, func(config *MultiAgentConfig) { + config.Routing.Middleware = []MiddlewareConfig{{Type: "cors", Enabled: true}} + config.Shared = &SharedConfig{Security: &SecurityConfig{CORS: &CORSConfig{ + Enabled: true, + AllowedOrigins: []string{"https://allowed.example"}, + AllowedMethods: []string{"GET", "POST"}, + MaxAge: 600, + }}} + }) + + server := httptest.NewServer(manager.GetRouter()) + defer server.Close() + + req, err := http.NewRequest(http.MethodPost, server.URL+"/solo", strings.NewReader(`{"input":"hi"}`)) + require.NoError(t, err) + req.Header.Set("Origin", "https://allowed.example") + resp, err := http.DefaultClient.Do(req) + require.NoError(t, err) + defer func() { _ = resp.Body.Close() }() + + assert.Equal(t, "https://allowed.example", resp.Header.Get("Access-Control-Allow-Origin")) + assert.Equal(t, "600", resp.Header.Get("Access-Control-Max-Age")) + + // A disallowed origin is not echoed back. + req2, err := http.NewRequest(http.MethodPost, server.URL+"/solo", strings.NewReader(`{"input":"hi"}`)) + require.NoError(t, err) + req2.Header.Set("Origin", "https://evil.example") + resp2, err := http.DefaultClient.Do(req2) + require.NoError(t, err) + defer func() { _ = resp2.Body.Close() }() + assert.Empty(t, resp2.Header.Get("Access-Control-Allow-Origin")) +} + +// Defect: the POST handler decoded whatever the client sent, with no ceiling on +// the body size. +func TestRegression_RequestBodyIsBounded(t *testing.T) { + manager, provider := newSingleAgentManager(t, nil) + server := httptest.NewServer(manager.GetRouter()) + defer server.Close() + + huge := `{"input":"` + strings.Repeat("A", MaxRequestBodyBytes+1024) + `"}` + resp, err := http.Post(server.URL+"/solo", "application/json", strings.NewReader(huge)) + require.NoError(t, err) + defer func() { _ = resp.Body.Close() }() + + assert.Equal(t, http.StatusBadRequest, resp.StatusCode) + assert.Equal(t, 0, provider.Calls(), "an over-sized body must never reach the agent") +} + +// --------------------------------------------------------------------------- +// Observability defects +// --------------------------------------------------------------------------- + +// Defect: /health hard-coded "status": "healthy" and HTTP 200 while listing +// agents whose health said otherwise, so a liveness probe pointed at it always +// passed. +func TestRegression_HealthEndpointReportsRealStatus(t *testing.T) { + manager, _ := newSingleAgentManager(t, nil) + server := httptest.NewServer(manager.GetRouter()) + defer server.Close() + + manager.updateAgentHealthStatus("solo", "healthy") + status, body := getJSON(t, server.URL+"/health") + assert.Equal(t, http.StatusOK, status) + assert.Equal(t, "healthy", body["status"]) + + manager.updateAgentHealthStatus("solo", "unhealthy") + status, body = getJSON(t, server.URL+"/health") + assert.Equal(t, http.StatusServiceUnavailable, status, "an unhealthy agent must not answer 200") + assert.Equal(t, "unhealthy", body["status"]) + + status, body = getJSON(t, server.URL+"/health/solo") + assert.Equal(t, http.StatusServiceUnavailable, status) + assert.Equal(t, "unhealthy", body["status"]) +} + +// Defect: the health check asked agent.IsRunning(), which reports whether the +// agent is *mid-execution*. An idle, perfectly healthy agent was therefore +// marked unhealthy on every single tick. +func TestRegression_IdleAgentIsHealthy(t *testing.T) { + provider := fakes.NewProvider("fake", "hi") + config := &MultiAgentConfig{ + Name: "health-checks", + Agents: map[string]*AgentConfig{"solo": newTestAgentConfig("solo", "fake")}, + Routing: &RoutingConfig{ + Type: "path", + Rules: []RoutingRule{{ID: "r", Pattern: "/solo", AgentID: "solo", Method: "POST"}}, + }, + Deployment: &DeploymentConfig{ + Type: "docker", Replicas: 1, + HealthCheck: &HealthCheckConfig{Enabled: true, PeriodSeconds: 1, FailureThreshold: 2}, + }, + } + + manager := newManagerWithProviders(t, config, map[string]*fakes.Provider{"fake": provider}) + + // The agent is idle - which used to be read as "not running" and therefore + // "unhealthy". + require.False(t, manager.agents["solo"].IsRunning()) + + results := manager.CheckHealthNow(context.Background()) + require.Contains(t, results, "solo") + assert.Equal(t, "healthy", results["solo"].Status) + assert.Equal(t, 0, results["solo"].ConsecutiveFails) + + state := manager.GetDeploymentState() + assert.Equal(t, "healthy", state.AgentStates["solo"].HealthStatus) + + // And an actually broken provider is detected. + provider.SetHealthy(false) + results = manager.CheckHealthNow(context.Background()) + assert.Equal(t, "unhealthy", results["solo"].Status) + assert.Equal(t, 1, results["solo"].ConsecutiveFails) + assert.Contains(t, results["solo"].LastError, "unhealthy") + + // Below the failure threshold the agent is degraded, not yet unhealthy. + assert.Equal(t, "degraded", manager.GetDeploymentState().AgentStates["solo"].HealthStatus) + + results = manager.CheckHealthNow(context.Background()) + assert.Equal(t, 2, results["solo"].ConsecutiveFails) + assert.Equal(t, "unhealthy", manager.GetDeploymentState().AgentStates["solo"].HealthStatus) +} + +// Defect: recordMetrics did all of its work inside "if the agent has a metrics +// entry", so a request routed to a missing agent - the 404 path, which is +// exactly a failure worth counting - incremented nothing. FailedRoutes was +// declared, serialized and never written at all. +func TestRegression_FailuresForUnknownAgentsAreCounted(t *testing.T) { + manager, _ := newSingleAgentManager(t, nil) + + before := manager.GetMetrics() + assert.Equal(t, int64(0), before.RoutingMetrics.FailedRoutes) + + handler := manager.createAgentHandler("does-not-exist", false) + recorder := httptest.NewRecorder() + handler(recorder, httptest.NewRequest(http.MethodPost, "/does-not-exist", strings.NewReader(`{"input":"hi"}`))) + + assert.Equal(t, http.StatusNotFound, recorder.Code) + + after := manager.GetMetrics() + assert.Equal(t, before.TotalErrors+1, after.TotalErrors, "a 404 for a missing agent is an error") + assert.Equal(t, int64(1), after.RoutingMetrics.FailedRoutes) + require.Contains(t, after.AgentMetrics, "does-not-exist") + assert.Equal(t, int64(1), after.AgentMetrics["does-not-exist"].ErrorCount) +} + +// Defect: DeploymentState.ErrorCount and LastError were declared, serialized to +// /deployment/status and never written, so the deployment always looked clean +// no matter how many executions failed. +func TestRegression_DeploymentLevelErrorsAreRecorded(t *testing.T) { + failing := fakes.NewProvider("failing", "").FailWith(errors.New("provider exploded")) + config := &MultiAgentConfig{ + Name: "deployment-errors", + Agents: map[string]*AgentConfig{"solo": newTestAgentConfig("solo", "failing")}, + Routing: &RoutingConfig{ + Type: "path", + Rules: []RoutingRule{{ID: "r", Pattern: "/solo", AgentID: "solo", Method: "POST"}}, + }, + } + + manager := newManagerWithProviders(t, config, map[string]*fakes.Provider{"failing": failing}) + server := httptest.NewServer(manager.GetRouter()) + defer server.Close() + + status, _ := postInput(t, server.URL+"/solo", "hi") + assert.Equal(t, http.StatusInternalServerError, status) + + state := manager.GetDeploymentState() + assert.Equal(t, 1, state.ErrorCount, "the deployment must count the failure") + assert.Contains(t, state.LastError, "provider exploded") + assert.Equal(t, int64(1), state.AgentStates["solo"].ErrorCount) + assert.Contains(t, state.AgentStates["solo"].LastError, "provider exploded") +} + +// Defect: GetDeploymentState dereferenced the struct and returned the copy, +// which still shared the AgentStates map and every *AgentState in it - so the +// "snapshot" kept changing under the caller and reading it raced with request +// handlers. +func TestRegression_GetDeploymentStateIsADeepCopy(t *testing.T) { + manager, _ := newSingleAgentManager(t, nil) + + snapshot := manager.GetDeploymentState() + require.Contains(t, snapshot.AgentStates, "solo") + assert.Equal(t, int64(0), snapshot.AgentStates["solo"].RequestCount) + + manager.updateAgentSuccess("solo") + manager.updateAgentError("solo", errors.New("boom")) + + assert.Equal(t, int64(0), snapshot.AgentStates["solo"].RequestCount, "the snapshot must not change") + assert.Equal(t, int64(0), snapshot.AgentStates["solo"].ErrorCount) + assert.Equal(t, 0, snapshot.ErrorCount) + + // Mutating the snapshot must not reach the manager either. + snapshot.AgentStates["solo"].Status = "tampered" + delete(snapshot.AgentStates, "solo") + fresh := manager.GetDeploymentState() + require.Contains(t, fresh.AgentStates, "solo") + assert.NotEqual(t, "tampered", fresh.AgentStates["solo"].Status) + assert.Equal(t, int64(1), fresh.AgentStates["solo"].RequestCount) +} + +// Defect: POST /deployment/restart logged "Restart requested", answered +// "restart_initiated" and did nothing at all. +func TestRegression_RestartActuallyRestartsAgents(t *testing.T) { + manager, _ := newSingleAgentManager(t, nil) + require.NoError(t, manager.Start(context.Background())) + + original := manager.agents["solo"] + manager.updateAgentError("solo", errors.New("earlier failure")) + require.Equal(t, 1, manager.GetDeploymentState().ErrorCount) + + server := httptest.NewServer(manager.GetRouter()) + defer server.Close() + + resp, err := http.Post(server.URL+"/deployment/restart", "application/json", nil) + require.NoError(t, err) + defer func() { _ = resp.Body.Close() }() + assert.Equal(t, http.StatusOK, resp.StatusCode) + + var body map[string]interface{} + require.NoError(t, json.NewDecoder(resp.Body).Decode(&body)) + assert.Equal(t, "restarted", body["status"]) + + state := manager.GetDeploymentState() + assert.Equal(t, "running", state.Status) + assert.Equal(t, 0, state.ErrorCount, "restart must clear the deployment error state") + assert.Empty(t, state.LastError) + assert.Equal(t, int64(0), state.AgentStates["solo"].ErrorCount) + assert.NotSame(t, original, manager.agents["solo"], "the agent must actually be rebuilt") + + // The manager is still usable afterwards. + status, _ := postInput(t, server.URL+"/solo", "hi") + assert.Equal(t, http.StatusOK, status) +} + +// --------------------------------------------------------------------------- +// Execution defects +// --------------------------------------------------------------------------- + +// Defect: the handler hard-coded a five minute deadline and never looked at +// AgentConfig.Timeout, so a config promising a short budget still let a slow +// provider hold the request open for minutes. +func TestRegression_AgentTimeoutComesFromConfig(t *testing.T) { + slow := fakes.NewProvider("slow", "eventually").WithDelay(5 * time.Second) + config := &MultiAgentConfig{ + Name: "timeout", + Agents: map[string]*AgentConfig{ + "solo": func() *AgentConfig { + c := newTestAgentConfig("solo", "slow") + c.Timeout = 150 * time.Millisecond + return c + }(), + }, + Routing: &RoutingConfig{ + Type: "path", + Rules: []RoutingRule{{ID: "r", Pattern: "/solo", AgentID: "solo", Method: "POST"}}, + }, + } + + manager := newManagerWithProviders(t, config, map[string]*fakes.Provider{"slow": slow}) + server := httptest.NewServer(manager.GetRouter()) + defer server.Close() + + start := time.Now() + status, _ := postInput(t, server.URL+"/solo", "hi") + elapsed := time.Since(start) + + assert.Equal(t, http.StatusGatewayTimeout, status, "a timeout must not be reported as a generic 500") + assert.Less(t, elapsed, 2*time.Second, "the configured 150ms budget must be applied, took %v", elapsed) +} + +// Defect: GetEnabledAgents returned every agent with a comment admitting it did +// no filtering, so a config that switched an agent off still got it created, +// health checked and routed to. +func TestRegression_DisabledAgentsAreNotCreated(t *testing.T) { + config := &MultiAgentConfig{ + Name: "disabled-agents", + Agents: map[string]*AgentConfig{ + "live": newTestAgentConfig("live", "fake"), + "off": func() *AgentConfig { + c := newTestAgentConfig("off", "fake") + c.Metadata = map[string]interface{}{"disabled": true} + return c + }(), + "off-string": func() *AgentConfig { + c := newTestAgentConfig("off-string", "fake") + c.Metadata = map[string]interface{}{"enabled": "false"} + return c + }(), + }, + Routing: &RoutingConfig{ + Type: "path", + Rules: []RoutingRule{{ID: "r", Pattern: "/live", AgentID: "live", Method: "POST"}}, + }, + Deployment: &DeploymentConfig{ + Type: "docker", Replicas: 1, + HealthCheck: &HealthCheckConfig{Enabled: true, PeriodSeconds: 1}, + }, + } + + assert.True(t, config.IsAgentEnabled("live")) + assert.False(t, config.IsAgentEnabled("off")) + assert.False(t, config.IsAgentEnabled("off-string")) + assert.Len(t, config.GetEnabledAgents(), 1) + + manager := newManagerWithProviders(t, config, map[string]*fakes.Provider{"fake": fakes.NewProvider("fake", "hi")}) + + state := manager.GetDeploymentState() + assert.Contains(t, state.AgentStates, "live") + assert.NotContains(t, state.AgentStates, "off") + assert.NotContains(t, state.AgentStates, "off-string") + + _, hasChecker := manager.HealthCheckerStatus("off") + assert.False(t, hasChecker, "a disabled agent must not be health checked") + _, hasChecker = manager.HealthCheckerStatus("live") + assert.True(t, hasChecker) +} + +// A routing rule pointing at a disabled agent is a config error rather than a +// route that silently 404s at runtime. +func TestRegression_RoutingToDisabledAgentIsRejected(t *testing.T) { + config := &MultiAgentConfig{ + Name: "route-to-disabled", + Agents: map[string]*AgentConfig{ + "off": func() *AgentConfig { + c := newTestAgentConfig("off", "fake") + c.Metadata = map[string]interface{}{"disabled": true} + return c + }(), + }, + Routing: &RoutingConfig{ + Type: "path", + Rules: []RoutingRule{{ID: "r", Pattern: "/off", AgentID: "off"}}, + }, + } + + err := config.Validate() + + require.Error(t, err) + assert.Contains(t, err.Error(), "is disabled") +} diff --git a/pkg/agent/multi_agent_lifecycle_test.go b/pkg/agent/multi_agent_lifecycle_test.go new file mode 100644 index 0000000..6ed7eea --- /dev/null +++ b/pkg/agent/multi_agent_lifecycle_test.go @@ -0,0 +1,861 @@ +// Copyright (c) 2024 GoLangGraph Team +// +// Licensed under the MIT License. See LICENSE file in the project root for full license information. + +package agent + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "runtime" + "strings" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/UnicoLab/GoLangGraph/pkg/llm" + "github.com/UnicoLab/GoLangGraph/pkg/tools" + "github.com/UnicoLab/GoLangGraph/test/fakes" +) + +// healthCheckedConfig builds a config whose agents are health checked, which is +// what starts the manager's background goroutines. +func healthCheckedConfig(name string, agentIDs []string, period int) *MultiAgentConfig { + config := &MultiAgentConfig{ + Name: name, + Version: "1.0", + Agents: make(map[string]*AgentConfig, len(agentIDs)), + Routing: &RoutingConfig{Type: "path"}, + Deployment: &DeploymentConfig{ + Type: "docker", + Replicas: 1, + HealthCheck: &HealthCheckConfig{ + Enabled: true, + PeriodSeconds: period, + TimeoutSeconds: 1, + FailureThreshold: 2, + }, + }, + } + for _, id := range agentIDs { + config.Agents[id] = newTestAgentConfig(id, "fake") + config.Routing.Rules = append(config.Routing.Rules, RoutingRule{ + ID: id, Pattern: "/" + id, AgentID: id, Method: "POST", + }) + } + return config +} + +// settleGoroutines waits for the goroutine count to come back down to baseline. +func settleGoroutines(t *testing.T, baseline int, timeout time.Duration) int { + t.Helper() + deadline := time.Now().Add(timeout) + current := runtime.NumGoroutine() + for time.Now().Before(deadline) { + current = runtime.NumGoroutine() + if current <= baseline { + return current + } + time.Sleep(20 * time.Millisecond) + } + return current +} + +// --------------------------------------------------------------------------- +// Lifecycle +// --------------------------------------------------------------------------- + +// Defect: setupHealthChecking launched one bare `go mam.runHealthChecker(...)` +// per agent from the constructor, looping on `for range ticker.C` with no +// cancellation. Nothing could ever stop them, so every manager ever built +// leaked a goroutine per agent for the lifetime of the process, and Stop +// returned "stopped" while they kept ticking. +func TestMultiAgentManagerStopReclaimsHealthCheckerGoroutines(t *testing.T) { + // Let anything left over from earlier tests wind down first. + baseline := settleGoroutines(t, 0, 2*time.Second) + + const agentCount = 6 + agentIDs := make([]string, 0, agentCount) + for i := 0; i < agentCount; i++ { + agentIDs = append(agentIDs, fmt.Sprintf("agent-%d", i)) + } + + llmManager := llm.NewProviderManager() + require.NoError(t, llmManager.RegisterProvider("fake", fakes.NewProvider("fake", "hi"))) + + manager, err := NewMultiAgentManager(healthCheckedConfig("leak-check", agentIDs, 1), llmManager, tools.NewToolRegistry()) + require.NoError(t, err) + + // Construction alone must not start anything. + assert.LessOrEqual(t, runtime.NumGoroutine(), baseline+2, + "constructing a manager must not spawn background goroutines") + + require.NoError(t, manager.Start(context.Background())) + + // The checkers really are running now. + require.Eventually(t, func() bool { + return runtime.NumGoroutine() >= baseline+agentCount + }, 3*time.Second, 20*time.Millisecond, "health checkers should be running after Start") + + stopCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + require.NoError(t, manager.Stop(stopCtx)) + + after := settleGoroutines(t, baseline, 5*time.Second) + assert.LessOrEqual(t, after, baseline, + "Stop must reclaim every health checker goroutine (baseline %d, after %d)", baseline, after) +} + +// Defect: runHealthChecker began with an unconditional +// time.Sleep(initial_delay_seconds). The shipped default config asks for 30s, +// so Stop had to wait out the whole delay before the goroutine so much as +// looked at cancellation - and there was no cancellation to look at. +func TestMultiAgentManagerStopIsPromptDespiteInitialDelay(t *testing.T) { + config := healthCheckedConfig("slow-start", []string{"solo"}, 1) + config.Deployment.HealthCheck.InitialDelaySeconds = 3600 + + llmManager := llm.NewProviderManager() + require.NoError(t, llmManager.RegisterProvider("fake", fakes.NewProvider("fake", "hi"))) + manager, err := NewMultiAgentManager(config, llmManager, tools.NewToolRegistry()) + require.NoError(t, err) + + require.NoError(t, manager.Start(context.Background())) + + stopCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + start := time.Now() + require.NoError(t, manager.Stop(stopCtx)) + elapsed := time.Since(start) + + assert.Less(t, elapsed, 2*time.Second, "Stop waited out the initial delay (%v)", elapsed) +} + +// Defect: an enabled health check with no period_seconds handed 0 to +// time.NewTicker, which panics - on a background goroutine, so the panic was +// unrecoverable and took the whole process down at startup. +func TestRegression_HealthCheckWithoutPeriodDoesNotPanic(t *testing.T) { + config := &MultiAgentConfig{ + Name: "no-period", + Agents: map[string]*AgentConfig{"solo": newTestAgentConfig("solo", "fake")}, + Routing: &RoutingConfig{Type: "path"}, + Deployment: &DeploymentConfig{ + Type: "docker", Replicas: 1, + HealthCheck: &HealthCheckConfig{Enabled: true}, // period_seconds omitted + }, + } + + manager := newManagerWithProviders(t, config, map[string]*fakes.Provider{"fake": fakes.NewProvider("fake", "hi")}) + + // A default period is substituted rather than passed straight to NewTicker. + assert.Equal(t, DefaultHealthCheckPeriod, config.Deployment.HealthCheck.Period()) + + require.NoError(t, manager.Start(context.Background())) + time.Sleep(200 * time.Millisecond) // long enough for a panicking goroutine to take us down + + stopCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + require.NoError(t, manager.Stop(stopCtx)) +} + +func TestMultiAgentManagerStopAndStartAreIdempotent(t *testing.T) { + llmManager := llm.NewProviderManager() + require.NoError(t, llmManager.RegisterProvider("fake", fakes.NewProvider("fake", "hi"))) + manager, err := NewMultiAgentManager(healthCheckedConfig("idempotent", []string{"a", "b"}, 1), llmManager, tools.NewToolRegistry()) + require.NoError(t, err) + + ctx := context.Background() + + // Stop before Start must be a no-op, not a panic or a hang. + require.NoError(t, manager.Stop(ctx)) + + require.NoError(t, manager.Start(ctx)) + require.NoError(t, manager.Start(ctx)) + assert.Equal(t, "running", manager.GetDeploymentState().Status) + + require.NoError(t, manager.Stop(ctx)) + require.NoError(t, manager.Stop(ctx)) + assert.Equal(t, "stopped", manager.GetDeploymentState().Status) + + // And it can come back up. + require.NoError(t, manager.Start(ctx)) + assert.Equal(t, "running", manager.GetDeploymentState().Status) + require.NoError(t, manager.Stop(ctx)) +} + +// Defect: Start and Stop took a context and ignored it entirely, so a caller +// whose context had already been canceled still got "started successfully". +func TestMultiAgentManagerStartHonorsContext(t *testing.T) { + manager, _ := newSingleAgentManager(t, nil) + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + err := manager.Start(ctx) + + require.Error(t, err) + assert.ErrorIs(t, err, context.Canceled) + assert.NotEqual(t, "running", manager.GetDeploymentState().Status) +} + +// --------------------------------------------------------------------------- +// Concurrency +// --------------------------------------------------------------------------- + +// Defect: handleDeploymentStatus and handleAgentStatus copied the deployment +// state shallowly (or kept the live *AgentState), released the lock and then +// let the JSON encoder walk structures that request handlers were concurrently +// writing. -race reported the write in updateAgentError/updateAgentSuccess +// against the read in the encoder. +// +// This test hammers agent execution and every management endpoint at once; it +// fails under -race against the unfixed code. +func TestMultiAgentConcurrentExecutionAndIntrospectionIsRaceFree(t *testing.T) { + const agentCount = 4 + + agentIDs := make([]string, 0, agentCount) + providers := map[string]*fakes.Provider{} + config := &MultiAgentConfig{ + Name: "concurrent", + Version: "1.0", + Agents: map[string]*AgentConfig{}, + Routing: &RoutingConfig{Type: "path"}, + } + for i := 0; i < agentCount; i++ { + id := fmt.Sprintf("agent-%d", i) + agentIDs = append(agentIDs, id) + providers[id] = fakes.NewProvider(id, "reply from "+id) + config.Agents[id] = newTestAgentConfig(id, id) + config.Routing.Rules = append(config.Routing.Rules, RoutingRule{ + ID: id, Pattern: "/" + id, AgentID: id, Method: "POST", + }) + } + + manager := newManagerWithProviders(t, config, providers) + require.NoError(t, manager.Start(context.Background())) + + server := httptest.NewServer(manager.GetRouter()) + defer server.Close() + + var wg sync.WaitGroup + var executed, introspected atomic.Int64 + + // Executors: one goroutine per agent so that Agent.Execute's own + // single-flight guard does not turn every call into an error. + for _, id := range agentIDs { + agentID := id + wg.Add(1) + go func() { + defer wg.Done() + for i := 0; i < 25; i++ { + status, _ := postInput(t, server.URL+"/"+agentID, "hello") + if status == http.StatusOK { + executed.Add(1) + } + } + }() + } + + // Introspectors: every read-only endpoint, plus the in-process accessors. + endpoints := []string{"/health", "/metrics", "/agents", "/config", "/routing", "/deployment/status", "/agents/agent-0", "/agents/agent-0/status", "/health/agent-1"} + for i := 0; i < 4; i++ { + wg.Add(1) + go func() { + defer wg.Done() + for j := 0; j < 30; j++ { + endpoint := endpoints[j%len(endpoints)] + resp, err := http.Get(server.URL + endpoint) + if err != nil { + continue + } + _ = resp.Body.Close() + introspected.Add(1) + + _ = manager.GetDeploymentState() + _ = manager.GetMetrics() + _, _ = manager.OverallHealth() + } + }() + } + + // Health checks run concurrently with everything else. + wg.Add(1) + go func() { + defer wg.Done() + for i := 0; i < 20; i++ { + manager.CheckHealthNow(context.Background()) + } + }() + + wg.Wait() + + assert.Equal(t, int64(agentCount*25), executed.Load(), "every request to a distinct agent should succeed") + assert.Positive(t, introspected.Load()) + + metrics := manager.GetMetrics() + for _, id := range agentIDs { + require.Contains(t, metrics.AgentMetrics, id) + assert.Equal(t, int64(25), metrics.AgentMetrics[id].RequestCount) + assert.Zero(t, metrics.AgentMetrics[id].ErrorCount) + } + assert.Zero(t, metrics.TotalErrors) +} + +// A failing agent must not contaminate its neighbours: the others keep serving +// and only the broken one accumulates errors. +func TestMultiAgentPartialFailureIsIsolatedAndReported(t *testing.T) { + good1 := fakes.NewProvider("good1", "fine") + good2 := fakes.NewProvider("good2", "fine") + // FailWith supplies an error for each of the first N calls; the fake keeps + // its own call counter, so listing the error five times fails five calls. + broken := fakes.NewProvider("broken", "").FailWith( + errors.New("upstream is down"), errors.New("upstream is down"), + errors.New("upstream is down"), errors.New("upstream is down"), + errors.New("upstream is down"), + ) + + config := &MultiAgentConfig{ + Name: "partial-failure", + Agents: map[string]*AgentConfig{ + "good1": newTestAgentConfig("good1", "good1"), + "good2": newTestAgentConfig("good2", "good2"), + "broken": newTestAgentConfig("broken", "broken"), + }, + Routing: &RoutingConfig{ + Type: "path", + Rules: []RoutingRule{ + {ID: "good1", Pattern: "/good1", AgentID: "good1", Method: "POST"}, + {ID: "good2", Pattern: "/good2", AgentID: "good2", Method: "POST"}, + {ID: "broken", Pattern: "/broken", AgentID: "broken", Method: "POST"}, + }, + }, + } + + manager := newManagerWithProviders(t, config, map[string]*fakes.Provider{ + "good1": good1, "good2": good2, "broken": broken, + }) + server := httptest.NewServer(manager.GetRouter()) + defer server.Close() + + var wg sync.WaitGroup + statuses := make(map[string][]int, 3) + var mu sync.Mutex + + for _, agentID := range []string{"good1", "good2", "broken"} { + id := agentID + wg.Add(1) + go func() { + defer wg.Done() + for i := 0; i < 5; i++ { + status, _ := postInput(t, server.URL+"/"+id, "work") + mu.Lock() + statuses[id] = append(statuses[id], status) + mu.Unlock() + } + }() + } + wg.Wait() + + for _, id := range []string{"good1", "good2"} { + for _, status := range statuses[id] { + assert.Equal(t, http.StatusOK, status, "healthy agent %s must keep serving", id) + } + } + for _, status := range statuses["broken"] { + assert.Equal(t, http.StatusInternalServerError, status) + } + + metrics := manager.GetMetrics() + assert.Zero(t, metrics.AgentMetrics["good1"].ErrorCount) + assert.Zero(t, metrics.AgentMetrics["good2"].ErrorCount) + assert.Equal(t, int64(5), metrics.AgentMetrics["broken"].ErrorCount) + assert.Equal(t, int64(5), metrics.TotalErrors, "the partial failure must not be reported as success") + + state := manager.GetDeploymentState() + assert.Equal(t, int64(0), state.AgentStates["good1"].ErrorCount) + assert.Equal(t, int64(5), state.AgentStates["broken"].ErrorCount) + assert.Contains(t, state.AgentStates["broken"].LastError, "upstream is down") + assert.Equal(t, 5, state.ErrorCount) +} + +// Cancelling the client's request must reach the agent and its provider rather +// than leaving the execution running to completion in the background. +func TestMultiAgentCancellationPropagatesToTheAgent(t *testing.T) { + slow := fakes.NewProvider("slow", "too late").WithDelay(10 * time.Second) + config := &MultiAgentConfig{ + Name: "cancellation", + Agents: map[string]*AgentConfig{"solo": newTestAgentConfig("solo", "slow")}, + Routing: &RoutingConfig{ + Type: "path", + Rules: []RoutingRule{{ID: "r", Pattern: "/solo", AgentID: "solo", Method: "POST"}}, + }, + } + + manager := newManagerWithProviders(t, config, map[string]*fakes.Provider{"slow": slow}) + server := httptest.NewServer(manager.GetRouter()) + defer server.Close() + + ctx, cancel := context.WithCancel(context.Background()) + req, err := http.NewRequestWithContext(ctx, http.MethodPost, server.URL+"/solo", strings.NewReader(`{"input":"hi"}`)) + require.NoError(t, err) + + done := make(chan error, 1) + go func() { + resp, reqErr := http.DefaultClient.Do(req) + if resp != nil { + _ = resp.Body.Close() + } + done <- reqErr + }() + + // Wait until the agent is actually executing, then cancel. + agent := manager.agents["solo"] + require.Eventually(t, agent.IsRunning, 3*time.Second, 10*time.Millisecond, "the agent should be executing") + cancel() + + select { + case reqErr := <-done: + require.Error(t, reqErr, "a canceled request must not return a normal response") + case <-time.After(3 * time.Second): + t.Fatal("the client request did not unblock after cancellation") + } + + // The crucial assertion: the run is abandoned well before the provider's + // 10s delay would have elapsed. + require.Eventually(t, func() bool { return !agent.IsRunning() }, 3*time.Second, 10*time.Millisecond, + "cancellation must reach the agent instead of letting it run to completion") + + assert.Equal(t, int64(1), manager.GetMetrics().AgentMetrics["solo"].ErrorCount) +} + +// Restart while requests are in flight must neither deadlock nor race. The +// restart path takes the manager lock, the lifecycle lock and the metrics lock +// in sequence while handlers hold them too. +func TestMultiAgentRestartUnderConcurrentTraffic(t *testing.T) { + providers := map[string]*fakes.Provider{ + "p0": fakes.NewProvider("p0", "a"), + "p1": fakes.NewProvider("p1", "b"), + } + config := healthCheckedConfig("restart-under-load", []string{"a0", "a1"}, 1) + config.Agents["a0"].Provider = "p0" + config.Agents["a1"].Provider = "p1" + + manager := newManagerWithProviders(t, config, providers) + require.NoError(t, manager.Start(context.Background())) + + server := httptest.NewServer(manager.GetRouter()) + defer server.Close() + + stop := make(chan struct{}) + var wg sync.WaitGroup + + for _, id := range []string{"a0", "a1"} { + agentID := id + wg.Add(1) + go func() { + defer wg.Done() + for { + select { + case <-stop: + return + default: + } + resp, err := http.Post(server.URL+"/"+agentID, "application/json", strings.NewReader(`{"input":"hi"}`)) + if err == nil { + _ = resp.Body.Close() + } + } + }() + } + + for i := 0; i < 3; i++ { + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + err := manager.Restart(ctx) + cancel() + require.NoError(t, err, "restart %d failed", i) + } + + close(stop) + wg.Wait() + + assert.Equal(t, "running", manager.GetDeploymentState().Status) + + // Still serving after all of that. + status, _ := postInput(t, server.URL+"/a0", "hi") + assert.Equal(t, http.StatusOK, status) +} + +// --------------------------------------------------------------------------- +// Malformed configuration +// --------------------------------------------------------------------------- + +func TestMultiAgentMalformedConfigurationsAreRejected(t *testing.T) { + validAgent := func() map[string]*AgentConfig { + return map[string]*AgentConfig{"solo": newTestAgentConfig("solo", "fake")} + } + + tests := []struct { + name string + config *MultiAgentConfig + wantErr string + }{ + { + name: "nil config", + config: nil, + wantErr: "multi-agent config is required", + }, + { + name: "no name", + config: &MultiAgentConfig{Agents: validAgent()}, + wantErr: "name is required", + }, + { + name: "no agents", + config: &MultiAgentConfig{Name: "x", Agents: map[string]*AgentConfig{}}, + wantErr: "at least one agent must be defined", + }, + { + name: "nil agent config", + config: &MultiAgentConfig{Name: "x", Agents: map[string]*AgentConfig{"ghost": nil}}, + wantErr: "agent ghost: configuration is empty", + }, + { + name: "agent without provider", + config: &MultiAgentConfig{Name: "x", Agents: map[string]*AgentConfig{ + "solo": {ID: "solo", Name: "solo", Type: AgentTypeChat, Model: "m"}, + }}, + wantErr: "provider is required", + }, + { + name: "negative agent timeout", + config: &MultiAgentConfig{Name: "x", Agents: func() map[string]*AgentConfig { + agents := validAgent() + agents["solo"].Timeout = -time.Second + return agents + }()}, + wantErr: "timeout must not be negative", + }, + { + name: "default agent missing", + config: &MultiAgentConfig{Name: "x", Agents: validAgent(), + Routing: &RoutingConfig{Type: "path", DefaultAgent: "nobody"}}, + wantErr: "default agent nobody does not exist", + }, + { + name: "rule for unknown agent", + config: &MultiAgentConfig{Name: "x", Agents: validAgent(), + Routing: &RoutingConfig{Type: "path", Rules: []RoutingRule{{ID: "r", Pattern: "/x", AgentID: "nobody"}}}}, + wantErr: "agent nobody does not exist", + }, + { + name: "duplicate rule ids", + config: &MultiAgentConfig{Name: "x", Agents: validAgent(), + Routing: &RoutingConfig{Type: "path", Rules: []RoutingRule{ + {ID: "dupe", Pattern: "/a", AgentID: "solo"}, + {ID: "dupe", Pattern: "/b", AgentID: "solo"}, + }}}, + wantErr: "duplicate rule id", + }, + { + name: "unknown match mode", + config: &MultiAgentConfig{Name: "x", Agents: validAgent(), + Routing: &RoutingConfig{Type: "path", Rules: []RoutingRule{ + {ID: "r", Pattern: "/a", Match: "fuzzy", AgentID: "solo"}, + }}}, + wantErr: "unsupported match mode", + }, + { + name: "invalid regex pattern", + config: &MultiAgentConfig{Name: "x", Agents: validAgent(), + Routing: &RoutingConfig{Type: "path", Rules: []RoutingRule{ + {ID: "r", Pattern: "([", Match: MatchRegex, AgentID: "solo"}, + }}}, + wantErr: "invalid regex pattern", + }, + { + name: "invalid condition operator", + config: &MultiAgentConfig{Name: "x", Agents: validAgent(), + Routing: &RoutingConfig{Type: "path", Rules: []RoutingRule{ + {ID: "r", Pattern: "/a", AgentID: "solo", Conditions: []RoutingCondition{ + {Type: ConditionHeader, Key: "X", Value: "y", Operator: "approximately"}, + }}, + }}}, + wantErr: "unsupported operator", + }, + { + name: "header condition without key", + config: &MultiAgentConfig{Name: "x", Agents: validAgent(), + Routing: &RoutingConfig{Type: "path", Rules: []RoutingRule{ + {ID: "r", Pattern: "/a", AgentID: "solo", Conditions: []RoutingCondition{ + {Type: ConditionHeader, Value: "y"}, + }}, + }}}, + wantErr: "requires a key", + }, + { + name: "zero replicas", + config: &MultiAgentConfig{Name: "x", Agents: validAgent(), + Deployment: &DeploymentConfig{Type: "docker", Replicas: 0}}, + wantErr: "replicas must be at least 1", + }, + { + name: "conflicting ports", + config: &MultiAgentConfig{Name: "x", Agents: validAgent(), + Deployment: &DeploymentConfig{Type: "docker", Replicas: 1, Networking: &NetworkingConfig{ + Ports: []PortConfig{{Name: "http", Port: 8080}, {Name: "metrics", Port: 8080}}, + }}}, + wantErr: "assigned to both", + }, + { + name: "negative health check period", + config: &MultiAgentConfig{Name: "x", Agents: validAgent(), + Deployment: &DeploymentConfig{Type: "docker", Replicas: 1, + HealthCheck: &HealthCheckConfig{Enabled: true, PeriodSeconds: -5}}}, + wantErr: "period_seconds must not be negative", + }, + { + name: "health check for unknown agent", + config: &MultiAgentConfig{Name: "x", Agents: validAgent(), + Deployment: &DeploymentConfig{Type: "docker", Replicas: 1, + HealthCheck: &HealthCheckConfig{Enabled: true, PeriodSeconds: 5, + AgentSpecific: map[string]*HealthCheckConfig{"nobody": {Enabled: true}}}}}, + wantErr: "agent does not exist", + }, + { + name: "scaling bounds inverted", + config: &MultiAgentConfig{Name: "x", Agents: validAgent(), + Deployment: &DeploymentConfig{Type: "docker", Replicas: 1, + Scaling: &ScalingConfig{Enabled: true, MinReplicas: 5, MaxReplicas: 2}}}, + wantErr: "must not be below min_replicas", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + // Validate must report, never panic. + err := tt.config.Validate() + require.Error(t, err) + assert.Contains(t, err.Error(), tt.wantErr) + + // And the manager must refuse to build from it. + llmManager := llm.NewProviderManager() + require.NoError(t, llmManager.RegisterProvider("fake", fakes.NewProvider("fake", "hi"))) + manager, buildErr := NewMultiAgentManager(tt.config, llmManager, tools.NewToolRegistry()) + assert.Error(t, buildErr) + assert.Nil(t, manager) + }) + } +} + +func TestLoadMultiAgentConfigFromFile(t *testing.T) { + dir := t.TempDir() + + valid := ` +name: from-file +version: "1.0" +agents: + solo: + id: solo + name: Solo + type: chat + model: fake-model + provider: fake +routing: + type: path + default_agent: solo + rules: + - id: r + pattern: /solo + agent_id: solo + method: POST +` + + write := func(name, content string) string { + path := filepath.Join(dir, name) + require.NoError(t, os.WriteFile(path, []byte(content), 0o600)) + return path + } + + t.Run("yaml", func(t *testing.T) { + config, err := LoadMultiAgentConfigFromFile(write("good.yaml", valid)) + require.NoError(t, err) + assert.Equal(t, "from-file", config.Name) + assert.Contains(t, config.Agents, "solo") + assert.Equal(t, "solo", config.Routing.DefaultAgent) + }) + + t.Run("json", func(t *testing.T) { + config, err := LoadMultiAgentConfigFromFile(write("good.yaml", valid)) + require.NoError(t, err) + encoded, err := json.Marshal(config) + require.NoError(t, err) + + reloaded, err := LoadMultiAgentConfigFromFile(write("good.json", string(encoded))) + require.NoError(t, err) + assert.Equal(t, config.Name, reloaded.Name) + }) + + t.Run("empty agent entry is reported not panicked", func(t *testing.T) { + _, err := LoadMultiAgentConfigFromFile(write("ghost.yaml", "name: ghosts\nagents:\n ghost:\n")) + require.Error(t, err) + assert.Contains(t, err.Error(), "configuration is empty") + }) + + t.Run("missing file", func(t *testing.T) { + _, err := LoadMultiAgentConfigFromFile(filepath.Join(dir, "nope.yaml")) + require.Error(t, err) + assert.Contains(t, err.Error(), "failed to read config file") + }) + + t.Run("unsupported extension", func(t *testing.T) { + _, err := LoadMultiAgentConfigFromFile(write("config.toml", "name = \"x\"")) + require.Error(t, err) + assert.Contains(t, err.Error(), "unsupported config file format") + }) + + t.Run("malformed yaml", func(t *testing.T) { + _, err := LoadMultiAgentConfigFromFile(write("broken.yaml", "name: [unterminated")) + require.Error(t, err) + assert.Contains(t, err.Error(), "failed to parse YAML config") + }) + + t.Run("invalid configuration", func(t *testing.T) { + _, err := LoadMultiAgentConfigFromFile(write("invalid.yaml", "version: \"1.0\"\n")) + require.Error(t, err) + assert.Contains(t, err.Error(), "invalid configuration") + }) +} + +// --------------------------------------------------------------------------- +// Rate limiter internals +// --------------------------------------------------------------------------- + +func TestTokenBucketRefillsOverTime(t *testing.T) { + now := time.Now() + bucket := newTokenBucket(2, time.Second, 2, now) + + allowed, _ := bucket.allow(now) + assert.True(t, allowed) + allowed, _ = bucket.allow(now) + assert.True(t, allowed) + + allowed, retry := bucket.allow(now) + assert.False(t, allowed, "the third request in the same instant must be rejected") + assert.Positive(t, retry) + + // After a full period the budget is back. + allowed, _ = bucket.allow(now.Add(time.Second)) + assert.True(t, allowed) +} + +// Unbounded growth: a per-IP limiter keyed by a caller-controlled value keeps +// one bucket per source address forever unless it is swept. +func TestRateLimiterKeyedBucketsStayBounded(t *testing.T) { + clock := time.Now() + limiter := &rateLimiter{ + perIP: &rateLimitRule{requests: 10, period: time.Minute, burst: 10}, + skipPaths: map[string]bool{}, + buckets: map[string]*tokenBucket{}, + lastSeen: map[string]time.Time{}, + now: func() time.Time { return clock }, + } + + for i := 0; i < maxRateLimitKeys+500; i++ { + // Advance the clock so early keys age out of the sweep window. + clock = clock.Add(time.Millisecond) + allowed, _ := limiter.allowKeyed(fmt.Sprintf("ip:10.0.%d.%d", i/256, i%256), limiter.perIP) + require.True(t, allowed) + } + + limiter.mu.Lock() + size := len(limiter.buckets) + seen := len(limiter.lastSeen) + limiter.mu.Unlock() + + assert.LessOrEqual(t, size, maxRateLimitKeys, "per-key buckets must not grow without bound") + assert.Equal(t, size, seen, "the bucket and last-seen maps must stay in step") +} + +func TestRateLimiterRequiresConfiguredLimits(t *testing.T) { + config := &MultiAgentConfig{ + Name: "rate-limit-no-config", + Agents: map[string]*AgentConfig{"solo": newTestAgentConfig("solo", "fake")}, + Routing: &RoutingConfig{ + Type: "path", + Rules: []RoutingRule{{ID: "r", Pattern: "/solo", AgentID: "solo", Method: "POST"}}, + Middleware: []MiddlewareConfig{{Type: "rate_limit", Enabled: true}}, + }, + } + + llmManager := llm.NewProviderManager() + require.NoError(t, llmManager.RegisterProvider("fake", fakes.NewProvider("fake", "hi"))) + + _, err := NewMultiAgentManager(config, llmManager, tools.NewToolRegistry()) + + require.Error(t, err) + assert.Contains(t, err.Error(), "no limits are configured") +} + +// The inline middleware spelling used by the shipped example config must work. +func TestRateLimiterReadsInlineMiddlewareConfig(t *testing.T) { + config := &MultiAgentConfig{ + Name: "inline-rate-limit", + Agents: map[string]*AgentConfig{"solo": newTestAgentConfig("solo", "fake")}, + Routing: &RoutingConfig{ + Type: "path", + Rules: []RoutingRule{{ID: "r", Pattern: "/solo", AgentID: "solo", Method: "POST"}}, + Middleware: []MiddlewareConfig{{ + Type: "rate_limit", + Enabled: true, + Config: map[string]interface{}{"requests_per_minute": 2, "burst_limit": 2}, + }}, + }, + } + + manager := newManagerWithProviders(t, config, map[string]*fakes.Provider{"fake": fakes.NewProvider("fake", "hi")}) + server := httptest.NewServer(manager.GetRouter()) + defer server.Close() + + limited := 0 + for i := 0; i < 5; i++ { + if status, _ := postInput(t, server.URL+"/solo", "hi"); status == http.StatusTooManyRequests { + limited++ + } + } + assert.Equal(t, 3, limited) +} + +// --------------------------------------------------------------------------- +// Metrics snapshots +// --------------------------------------------------------------------------- + +func TestGetMetricsReturnsAnIndependentSnapshot(t *testing.T) { + manager, _ := newSingleAgentManager(t, nil) + + snapshot := manager.GetMetrics() + require.Contains(t, snapshot.AgentMetrics, "solo") + + manager.recordMetrics("solo", time.Millisecond, true) + manager.updateRoutingMetrics("solo", false) + manager.recordFailedRoute() + + assert.Equal(t, int64(0), snapshot.AgentMetrics["solo"].RequestCount, "the snapshot must not change") + assert.Equal(t, int64(0), snapshot.TotalErrors) + assert.Equal(t, int64(0), snapshot.RoutingMetrics.FailedRoutes) + + fresh := manager.GetMetrics() + assert.Equal(t, int64(1), fresh.AgentMetrics["solo"].RequestCount) + assert.Equal(t, int64(1), fresh.TotalErrors) + assert.Equal(t, int64(1), fresh.RoutingMetrics.RoutingDecisions["solo"]) + assert.Equal(t, int64(1), fresh.RoutingMetrics.FailedRoutes) +} diff --git a/pkg/agent/multi_agent_manager.go b/pkg/agent/multi_agent_manager.go index faa609b..535bf3c 100644 --- a/pkg/agent/multi_agent_manager.go +++ b/pkg/agent/multi_agent_manager.go @@ -8,11 +8,15 @@ package agent import ( "context" + "crypto/subtle" "encoding/json" + "errors" "fmt" + "net" "net/http" "os" "path/filepath" + "regexp" "sort" "strings" "sync" @@ -26,6 +30,14 @@ import ( "github.com/UnicoLab/GoLangGraph/pkg/tools" ) +// MaxRequestBodyBytes caps how much of a request body the agent handlers will +// read. Without a cap a single client could stream an unbounded body into the +// JSON decoder and exhaust the process's memory. +const MaxRequestBodyBytes = 1 << 20 // 1 MiB + +// DefaultAgentExecutionTimeout is used when an agent config sets no timeout. +const DefaultAgentExecutionTimeout = 5 * time.Minute + // MultiAgentManager manages multiple agents with routing and deployment capabilities type MultiAgentManager struct { config *MultiAgentConfig @@ -42,8 +54,21 @@ type MultiAgentManager struct { healthCheckers map[string]*HealthChecker healthMu sync.RWMutex + // Health checker lifecycle. The checkers used to be started from the + // constructor as bare goroutines with no cancellation and an endless + // "for range ticker.C" loop, so every manager ever built leaked one + // goroutine per agent for the lifetime of the process and Stop had no way + // to reclaim them. + lifecycleMu sync.Mutex + healthCancel context.CancelFunc + healthWG sync.WaitGroup + healthRunning bool + // Metrics and monitoring metrics *MultiAgentMetrics + + // Request limiting + limiter *rateLimiter } // MiddlewareFunc defines middleware function signature @@ -60,6 +85,29 @@ type DeploymentState struct { Metadata map[string]interface{} `json:"metadata"` } +// Clone returns a deep copy of the deployment state. +// +// GetDeploymentState and the /deployment/status handler used to dereference the +// struct while holding the read lock and hand the result out - but the copy +// shares the AgentStates map and every *AgentState in it, so callers (and the +// JSON encoder, after the lock was released) read fields that request handlers +// were concurrently writing. The race detector flags it on any concurrent load. +func (ds *DeploymentState) Clone() *DeploymentState { + if ds == nil { + return nil + } + clone := *ds + clone.AgentStates = make(map[string]*AgentState, len(ds.AgentStates)) + for id, state := range ds.AgentStates { + clone.AgentStates[id] = state.Clone() + } + clone.Metadata = make(map[string]interface{}, len(ds.Metadata)) + for k, v := range ds.Metadata { + clone.Metadata[k] = v + } + return &clone +} + // AgentState tracks the state of individual agents type AgentState struct { ID string `json:"id"` @@ -74,14 +122,84 @@ type AgentState struct { Metadata map[string]interface{} `json:"metadata"` } -// HealthChecker performs health checks for agents +// Clone returns a deep copy of the agent state. +func (as *AgentState) Clone() *AgentState { + if as == nil { + return nil + } + clone := *as + clone.Metadata = make(map[string]interface{}, len(as.Metadata)) + for k, v := range as.Metadata { + clone.Metadata[k] = v + } + return &clone +} + +// HealthChecker performs health checks for agents. +// +// The mutable fields are guarded by mu: they are written by the checker's own +// goroutine and read by callers of Snapshot (and by the /health handlers), so +// leaving them bare was a data race waiting for the first reader. type HealthChecker struct { - AgentID string - Config *HealthCheckConfig - LastCheck time.Time - Status string - ConsecutiveFails int - Logger *logrus.Logger + AgentID string + Config *HealthCheckConfig + Logger *logrus.Logger + + mu sync.RWMutex + lastCheck time.Time + status string + consecutiveFails int + lastError string +} + +// HealthCheckerSnapshot is an immutable view of a checker's state. +type HealthCheckerSnapshot struct { + AgentID string `json:"agent_id"` + Status string `json:"status"` + LastCheck time.Time `json:"last_check"` + ConsecutiveFails int `json:"consecutive_fails"` + LastError string `json:"last_error,omitempty"` +} + +// Snapshot returns the checker's current state under its lock. +func (hc *HealthChecker) Snapshot() HealthCheckerSnapshot { + hc.mu.RLock() + defer hc.mu.RUnlock() + return HealthCheckerSnapshot{ + AgentID: hc.AgentID, + Status: hc.status, + LastCheck: hc.lastCheck, + ConsecutiveFails: hc.consecutiveFails, + LastError: hc.lastError, + } +} + +// record stores the outcome of one check and reports the resulting state plus +// whether the status changed. +func (hc *HealthChecker) record(healthy bool, status string, err error) (HealthCheckerSnapshot, bool) { + hc.mu.Lock() + defer hc.mu.Unlock() + + previous := hc.status + hc.lastCheck = time.Now() + hc.status = status + if healthy { + hc.consecutiveFails = 0 + hc.lastError = "" + } else { + hc.consecutiveFails++ + if err != nil { + hc.lastError = err.Error() + } + } + + return HealthCheckerSnapshot{ + AgentID: hc.AgentID, + Status: hc.status, + LastCheck: hc.lastCheck, + ConsecutiveFails: hc.consecutiveFails, + LastError: hc.lastError, + }, previous != status } // MultiAgentMetrics tracks metrics for multi-agent system @@ -155,7 +273,8 @@ func NewMultiAgentManager(config *MultiAgentConfig, llmManager *llm.ProviderMana return nil, fmt.Errorf("failed to setup routing: %w", err) } - // Setup health checking + // Build the health checkers. They are only *started* by Start, so a manager + // that is constructed and dropped does not leave goroutines behind. if err := manager.setupHealthChecking(); err != nil { return nil, fmt.Errorf("failed to setup health checking: %w", err) } @@ -163,15 +282,35 @@ func NewMultiAgentManager(config *MultiAgentConfig, llmManager *llm.ProviderMana return manager, nil } -// initializeAgents creates and initializes all agents +// initializeAgents creates and initializes all agents. +// +// The agents, deployment state and per-agent metrics are rebuilt into fresh +// maps so the function is safe to run again from Restart. Agent construction +// happens outside the manager lock: it calls into user-supplied definitions and +// factories, and holding mam.mu across that invites a deadlock if any of them +// reaches back into the manager. func (mam *MultiAgentManager) initializeAgents() error { - mam.mu.Lock() - defer mam.mu.Unlock() - registry := GetGlobalRegistry() - for agentID, agentConfig := range mam.config.Agents { + agents := make(map[string]Agent) + states := make(map[string]*AgentState) + agentMetrics := make(map[string]*AgentMetrics) + + // Only enabled agents are created; a disabled agent must not be started, + // routed to or health checked. + enabled := mam.config.GetEnabledAgents() + agentIDs := make([]string, 0, len(enabled)) + for agentID := range enabled { + agentIDs = append(agentIDs, agentID) + } + sort.Strings(agentIDs) + + for _, agentID := range agentIDs { + agentConfig := enabled[agentID] // Ensure agent has an ID + if agentConfig == nil { + return fmt.Errorf("agent %s: configuration is empty", agentID) + } if agentConfig.ID == "" { agentConfig.ID = agentID } @@ -214,10 +353,14 @@ func (mam *MultiAgentManager) initializeAgents() error { } } - mam.agents[agentID] = agent + if agent == nil { + return fmt.Errorf("agent %s: builder returned no agent", agentID) + } + + agents[agentID] = agent // Initialize agent state - mam.deploymentState.AgentStates[agentID] = &AgentState{ + states[agentID] = &AgentState{ ID: agentID, Status: "initialized", StartedAt: time.Now(), @@ -229,7 +372,7 @@ func (mam *MultiAgentManager) initializeAgents() error { } // Initialize agent metrics - mam.metrics.AgentMetrics[agentID] = &AgentMetrics{ + agentMetrics[agentID] = &AgentMetrics{ RequestCount: 0, ErrorCount: 0, AverageLatency: 0, @@ -239,19 +382,40 @@ func (mam *MultiAgentManager) initializeAgents() error { mam.logger.WithField("agent_id", agentID).Info("Agent initialized") } + mam.mu.Lock() + mam.agents = agents + mam.deploymentState.AgentStates = states + mam.deploymentState.UpdatedAt = time.Now() + mam.mu.Unlock() + + mam.metrics.mu.Lock() + mam.metrics.AgentMetrics = agentMetrics + mam.metrics.LastUpdated = time.Now() + mam.metrics.mu.Unlock() + return nil } // setupRouting configures HTTP routing for multi-agent requests func (mam *MultiAgentManager) setupRouting() error { + // A config without a routing block is valid (Validate only checks routing + // when it is present), but this used to dereference it unconditionally and + // panic before the manager was ever returned. + routing := mam.config.Routing + if routing == nil { + routing = &RoutingConfig{} + } + // Setup global middleware - for _, middlewareConfig := range mam.config.Routing.Middleware { - if middlewareConfig.Enabled { - middleware := mam.createMiddleware(middlewareConfig) - if middleware != nil { - mam.middleware = append(mam.middleware, middleware) - } + for _, middlewareConfig := range routing.Middleware { + if !middlewareConfig.Enabled { + continue + } + middleware, err := mam.createMiddleware(middlewareConfig) + if err != nil { + return fmt.Errorf("middleware %q: %w", middlewareConfig.Type, err) } + mam.middleware = append(mam.middleware, middleware) } // Apply middleware to router @@ -262,73 +426,197 @@ func (mam *MultiAgentManager) setupRouting() error { // Add metrics middleware mam.router.Use(mux.MiddlewareFunc(mam.metricsMiddleware)) - // Sort routing rules by priority - rules := make([]RoutingRule, len(mam.config.Routing.Rules)) - copy(rules, mam.config.Routing.Rules) - sort.Slice(rules, func(i, j int) bool { - return rules[i].Priority > rules[j].Priority - }) - // Setup management endpoints FIRST so they don't get caught by other routes mam.setupManagementEndpoints() - // Setup routing rules - for _, rule := range rules { - mam.setupRoutingRule(rule) + // Setup routing rules, highest priority first. + for _, rule := range mam.config.SortedRules() { + if err := mam.setupRoutingRule(routing, rule); err != nil { + return fmt.Errorf("routing rule %q: %w", rule.ID, err) + } } // Setup default route if configured (this should be LAST) - if mam.config.Routing.DefaultAgent != "" { - mam.router.PathPrefix("/").HandlerFunc(mam.createAgentHandler(mam.config.Routing.DefaultAgent, true)) + if routing.DefaultAgent != "" { + mam.router.PathPrefix("/").HandlerFunc(mam.createAgentHandler(routing.DefaultAgent, true)) } return nil } -// setupRoutingRule sets up a single routing rule -func (mam *MultiAgentManager) setupRoutingRule(rule RoutingRule) { +// setupRoutingRule sets up a single routing rule. +// +// Every failure path here used to leave `route` nil and return silently, so a +// pattern the router could not express - a header pattern whose value contains +// a colon, for instance - produced a config that loaded cleanly and an agent +// that was simply unreachable. Failures are now reported and abort startup. +func (mam *MultiAgentManager) setupRoutingRule(routing *RoutingConfig, rule RoutingRule) error { handler := mam.createAgentHandler(rule.AgentID, false) + conditions, err := compileConditions(rule.Conditions) + if err != nil { + return err + } + var route *mux.Route - switch mam.config.Routing.Type { - case "path": - route = mam.router.Path(rule.Pattern) + switch strings.ToLower(strings.TrimSpace(routing.Type)) { + case "", "path": + matcher, err := pathMatcher(rule) + if err != nil { + return err + } + if matcher != nil { + route = mam.router.MatcherFunc(func(r *http.Request, rm *mux.RouteMatch) bool { + return matcher(r.URL.Path) + }) + } else { + route = mam.router.Path(rule.Pattern) + } case "host": route = mam.router.Host(rule.Pattern) case "header": - // Extract header key and value from pattern - parts := strings.Split(rule.Pattern, ":") - if len(parts) == 2 { - route = mam.router.MatcherFunc(func(r *http.Request, rm *mux.RouteMatch) bool { - return r.Header.Get(parts[0]) == parts[1] - }) + // SplitN keeps colons that belong to the value ("Authorization: Bearer + // x"); plain Split rejected any such pattern by producing three parts. + parts := strings.SplitN(rule.Pattern, ":", 2) + if len(parts) != 2 { + return fmt.Errorf("header pattern %q must be \"Header-Name: value\"", rule.Pattern) } + name := strings.TrimSpace(parts[0]) + value := strings.TrimSpace(parts[1]) + if name == "" { + return fmt.Errorf("header pattern %q has an empty header name", rule.Pattern) + } + route = mam.router.MatcherFunc(func(r *http.Request, rm *mux.RouteMatch) bool { + return r.Header.Get(name) == value + }) case "query": - // Extract query key and value from pattern - parts := strings.Split(rule.Pattern, "=") - if len(parts) == 2 { - route = mam.router.MatcherFunc(func(r *http.Request, rm *mux.RouteMatch) bool { - return r.URL.Query().Get(parts[0]) == parts[1] - }) + parts := strings.SplitN(rule.Pattern, "=", 2) + if len(parts) != 2 { + return fmt.Errorf("query pattern %q must be \"key=value\"", rule.Pattern) + } + key := strings.TrimSpace(parts[0]) + value := parts[1] + if key == "" { + return fmt.Errorf("query pattern %q has an empty key", rule.Pattern) } + route = mam.router.MatcherFunc(func(r *http.Request, rm *mux.RouteMatch) bool { + return r.URL.Query().Get(key) == value + }) default: - route = mam.router.Path(rule.Pattern) + return fmt.Errorf("unsupported routing type %q", routing.Type) + } + + if route == nil { + return fmt.Errorf("pattern %q could not be installed", rule.Pattern) + } + + // Conditions are part of the rule's match: a rule that declares them must + // not fire for a request that fails them. + if len(conditions) > 0 { + route = route.MatcherFunc(func(r *http.Request, rm *mux.RouteMatch) bool { + return matchConditions(r, conditions) + }) } - if route != nil { - if rule.Method != "" { - route = route.Methods(rule.Method) + if rule.Method != "" { + route = route.Methods(rule.Method) + } + route.Handler(handler) + + mam.logger.WithFields(logrus.Fields{ + "rule_id": rule.ID, + "pattern": rule.Pattern, + "match": rule.MatchMode(), + "agent_id": rule.AgentID, + "method": rule.Method, + "priority": rule.Priority, + "conditions": len(conditions), + }).Info("Routing rule configured") + + return nil +} + +// pathMatcher returns a path predicate for rules that need a match mode mux +// cannot express directly. It returns nil for plain exact matching so mux keeps +// handling its own path templates. +func pathMatcher(rule RoutingRule) (func(string) bool, error) { + switch rule.MatchMode() { + case MatchExact: + return nil, nil + case MatchPrefix: + return func(path string) bool { return strings.HasPrefix(path, rule.Pattern) }, nil + case MatchSuffix: + return func(path string) bool { return strings.HasSuffix(path, rule.Pattern) }, nil + case MatchContains: + return func(path string) bool { return strings.Contains(path, rule.Pattern) }, nil + case MatchRegex: + re, err := regexp.Compile(rule.Pattern) + if err != nil { + return nil, fmt.Errorf("invalid regex pattern %q: %w", rule.Pattern, err) } - route.Handler(handler) + return re.MatchString, nil + default: + return nil, fmt.Errorf("unsupported match mode %q", rule.Match) + } +} - mam.logger.WithFields(logrus.Fields{ - "rule_id": rule.ID, - "pattern": rule.Pattern, - "agent_id": rule.AgentID, - "method": rule.Method, - "priority": rule.Priority, - }).Info("Routing rule configured") +// compiledCondition pairs a condition with the request accessor it reads. +type compiledCondition struct { + condition RoutingCondition + value func(*http.Request) string +} + +// compileConditions resolves each routing condition to a request accessor, +// rejecting the ones the router cannot evaluate instead of ignoring them. +func compileConditions(conditions []RoutingCondition) ([]compiledCondition, error) { + compiled := make([]compiledCondition, 0, len(conditions)) + for _, cond := range conditions { + if err := validateCondition(cond); err != nil { + return nil, err + } + condition := cond + switch strings.ToLower(strings.TrimSpace(cond.Type)) { + case ConditionHeader: + compiled = append(compiled, compiledCondition{condition, func(r *http.Request) string { + return r.Header.Get(condition.Key) + }}) + case ConditionQuery: + compiled = append(compiled, compiledCondition{condition, func(r *http.Request) string { + return r.URL.Query().Get(condition.Key) + }}) + case ConditionIP: + compiled = append(compiled, compiledCondition{condition, clientIP}) + case ConditionMethod: + compiled = append(compiled, compiledCondition{condition, func(r *http.Request) string { + return r.Method + }}) + } + } + return compiled, nil +} + +func matchConditions(r *http.Request, conditions []compiledCondition) bool { + for _, compiled := range conditions { + if !compiled.condition.Evaluate(compiled.value(r)) { + return false + } } + return true +} + +// clientIP extracts the caller's address, preferring X-Forwarded-For when the +// server sits behind a proxy. +func clientIP(r *http.Request) string { + if forwarded := r.Header.Get("X-Forwarded-For"); forwarded != "" { + if first, _, found := strings.Cut(forwarded, ","); found { + return strings.TrimSpace(first) + } + return strings.TrimSpace(forwarded) + } + if host, _, err := net.SplitHostPort(r.RemoteAddr); err == nil { + return host + } + return r.RemoteAddr } // createAgentHandler creates an HTTP handler for a specific agent @@ -340,6 +628,8 @@ func (mam *MultiAgentManager) createAgentHandler(agentID string, isDefault bool) agent, exists := mam.getAgent(agentID) if !exists { mam.recordMetrics(agentID, time.Since(start), true) + mam.updateRoutingMetrics(agentID, isDefault) + mam.recordFailedRoute() http.Error(w, fmt.Sprintf("Agent %s not found", agentID), http.StatusNotFound) return } @@ -347,19 +637,28 @@ func (mam *MultiAgentManager) createAgentHandler(agentID string, isDefault bool) // Update routing metrics mam.updateRoutingMetrics(agentID, isDefault) + // Per-agent budgets can only be applied here, where the agent is known. + if allowed, retry := mam.limiter.allowAgent(agentID); !allowed { + mam.recordMetrics(agentID, time.Since(start), true) + writeRateLimited(w, retry) + return + } + // Parse request var input string switch r.Method { - case "GET": + case http.MethodGet: input = r.URL.Query().Get("input") if input == "" { input = r.URL.Query().Get("q") } - case "POST": + case http.MethodPost: var requestData struct { Input string `json:"input"` } - if err := json.NewDecoder(r.Body).Decode(&requestData); err != nil { + // Cap the body: the decoder used to read whatever the client sent. + decoder := json.NewDecoder(http.MaxBytesReader(w, r.Body, MaxRequestBodyBytes)) + if err := decoder.Decode(&requestData); err != nil { mam.recordMetrics(agentID, time.Since(start), true) http.Error(w, "Invalid request body", http.StatusBadRequest) return @@ -377,21 +676,39 @@ func (mam *MultiAgentManager) createAgentHandler(agentID string, isDefault bool) return } - // Execute agent - ctx, cancel := context.WithTimeout(r.Context(), 5*time.Minute) + // Execute the agent under the timeout its own config asks for. The + // handler used to hard-code five minutes and ignore AgentConfig.Timeout + // entirely, so a config that promised a 100ms budget still let a slow + // provider hold the request open for minutes. + ctx, cancel := context.WithTimeout(r.Context(), mam.executionTimeout(agent)) defer cancel() execution, err := agent.Execute(ctx, input) if err != nil { mam.recordMetrics(agentID, time.Since(start), true) mam.updateAgentError(agentID, err) - http.Error(w, fmt.Sprintf("Agent execution failed: %v", err), http.StatusInternalServerError) + http.Error(w, fmt.Sprintf("Agent execution failed: %v", err), executionErrorStatus(r.Context(), err)) return } - // Record successful metrics - mam.recordMetrics(agentID, time.Since(start), false) - mam.updateAgentSuccess(agentID) + // Defensive: Execute currently reports every failure through err, but + // an execution that comes back with Success=false is a failure whether + // or not err was set, and counting it as a success would understate the + // error rate. + if execution != nil && !execution.Success { + mam.recordMetrics(agentID, time.Since(start), true) + if execution.Error != nil { + mam.updateAgentError(agentID, execution.Error) + } else if execution.ErrorMessage != "" { + mam.updateAgentError(agentID, errors.New(execution.ErrorMessage)) + } else { + mam.updateAgentError(agentID, errors.New("agent execution did not succeed")) + } + } else { + // Record successful metrics + mam.recordMetrics(agentID, time.Since(start), false) + mam.updateAgentSuccess(agentID) + } // Return response w.Header().Set("Content-Type", "application/json") @@ -400,10 +717,36 @@ func (mam *MultiAgentManager) createAgentHandler(agentID string, isDefault bool) "execution": execution, "timestamp": time.Now(), } - json.NewEncoder(w).Encode(response) + _ = json.NewEncoder(w).Encode(response) } } +// executionTimeout returns the deadline to apply to one agent run. +func (mam *MultiAgentManager) executionTimeout(agent Agent) time.Duration { + if agent != nil { + if config := agent.GetConfig(); config != nil && config.Timeout > 0 { + return config.Timeout + } + } + return DefaultAgentExecutionTimeout +} + +// executionErrorStatus distinguishes "we ran out of time" from "the agent +// failed", which used to be reported identically as a 500. +func executionErrorStatus(requestCtx context.Context, err error) int { + if errors.Is(err, context.DeadlineExceeded) { + return http.StatusGatewayTimeout + } + if errors.Is(err, context.Canceled) { + if requestCtx.Err() != nil { + // The client went away; nothing will read the response anyway. + return http.StatusRequestTimeout + } + return http.StatusGatewayTimeout + } + return http.StatusInternalServerError +} + // setupManagementEndpoints sets up management and monitoring endpoints func (mam *MultiAgentManager) setupManagementEndpoints() { // Health check endpoint @@ -427,107 +770,289 @@ func (mam *MultiAgentManager) setupManagementEndpoints() { mam.router.HandleFunc("/deployment/restart", mam.handleRestart).Methods("POST") } -// setupHealthChecking initializes health checking for agents +// setupHealthChecking builds the per-agent health checkers. +// +// It deliberately does not start them: the goroutines belong to the manager's +// running lifetime and are launched by Start and reclaimed by Stop. func (mam *MultiAgentManager) setupHealthChecking() error { + mam.healthMu.Lock() + defer mam.healthMu.Unlock() + + mam.healthCheckers = make(map[string]*HealthChecker) + if mam.config.Deployment == nil || mam.config.Deployment.HealthCheck == nil || !mam.config.Deployment.HealthCheck.Enabled { return nil } - mam.healthMu.Lock() - defer mam.healthMu.Unlock() - - for agentID := range mam.config.Agents { + for agentID := range mam.config.GetEnabledAgents() { healthConfig := mam.config.Deployment.HealthCheck // Check for agent-specific health check config - if agentSpecific, exists := healthConfig.AgentSpecific[agentID]; exists { + if agentSpecific, exists := healthConfig.AgentSpecific[agentID]; exists && agentSpecific != nil { healthConfig = agentSpecific } - checker := &HealthChecker{ + mam.healthCheckers[agentID] = &HealthChecker{ AgentID: agentID, Config: healthConfig, - Status: "unknown", Logger: logrus.New(), + status: "unknown", } + } + + return nil +} - mam.healthCheckers[agentID] = checker +// startHealthCheckers launches one goroutine per checker, all cancellable. +func (mam *MultiAgentManager) startHealthCheckers() { + mam.lifecycleMu.Lock() + defer mam.lifecycleMu.Unlock() - // Start health checking goroutine - go mam.runHealthChecker(checker) + if mam.healthRunning { + return } - return nil + mam.healthMu.RLock() + checkers := make([]*HealthChecker, 0, len(mam.healthCheckers)) + for _, checker := range mam.healthCheckers { + checkers = append(checkers, checker) + } + mam.healthMu.RUnlock() + + if len(checkers) == 0 { + return + } + + ctx, cancel := context.WithCancel(context.Background()) + mam.healthCancel = cancel + mam.healthRunning = true + + for _, checker := range checkers { + mam.healthWG.Add(1) + go func(c *HealthChecker) { + defer mam.healthWG.Done() + mam.runHealthChecker(ctx, c) + }(checker) + } } -// runHealthChecker runs health checks for an agent -func (mam *MultiAgentManager) runHealthChecker(checker *HealthChecker) { - ticker := time.NewTicker(time.Duration(checker.Config.PeriodSeconds) * time.Second) - defer ticker.Stop() +// stopHealthCheckers cancels the checkers and waits for them to exit, giving up +// when ctx expires so a caller with a deadline is never blocked forever. +func (mam *MultiAgentManager) stopHealthCheckers(ctx context.Context) error { + mam.lifecycleMu.Lock() + if !mam.healthRunning { + mam.lifecycleMu.Unlock() + return nil + } + cancel := mam.healthCancel + mam.healthCancel = nil + mam.healthRunning = false + mam.lifecycleMu.Unlock() + + if cancel != nil { + cancel() + } - // Initial delay - time.Sleep(time.Duration(checker.Config.InitialDelaySeconds) * time.Second) + done := make(chan struct{}) + go func() { + mam.healthWG.Wait() + close(done) + }() - for range ticker.C { - mam.performHealthCheck(checker) + if ctx == nil { + <-done + return nil + } + + select { + case <-done: + return nil + case <-ctx.Done(): + return fmt.Errorf("timed out waiting for health checkers to stop: %w", ctx.Err()) } } -// performHealthCheck performs a single health check -func (mam *MultiAgentManager) performHealthCheck(checker *HealthChecker) { - checker.LastCheck = time.Now() +// runHealthChecker runs health checks for an agent until ctx is canceled. +func (mam *MultiAgentManager) runHealthChecker(ctx context.Context, checker *HealthChecker) { + // Config.Period substitutes a default for a missing period_seconds. The + // raw value was handed to time.NewTicker, and time.NewTicker(0) panics - + // on a background goroutine, so an enabled health check with no interval + // crashed the whole process at startup. + ticker := time.NewTicker(checker.Config.Period()) + defer ticker.Stop() + + // Initial delay, interruptible: an unconditional Sleep here meant Stop had + // to wait out initial_delay_seconds (30s in the shipped default config) + // before the goroutine would even look at cancellation. + if delay := checker.Config.InitialDelay(); delay > 0 { + timer := time.NewTimer(delay) + select { + case <-ctx.Done(): + timer.Stop() + return + case <-timer.C: + } + } + + for { + select { + case <-ctx.Done(): + return + case <-ticker.C: + mam.performHealthCheck(ctx, checker) + } + } +} - // Simple health check - in a real implementation, this would make HTTP requests +// performHealthCheck performs a single health check. +// +// The old check asked agent.IsRunning(), which reports whether the agent is +// *mid-execution*: an idle, perfectly healthy agent was therefore marked +// unhealthy on every tick and a busy one looked healthy. The real question is +// whether the agent exists and its LLM provider is reachable. +func (mam *MultiAgentManager) performHealthCheck(ctx context.Context, checker *HealthChecker) { agent, exists := mam.getAgent(checker.AgentID) if !exists { - checker.Status = "not_found" - checker.ConsecutiveFails++ + snapshot, changed := checker.record(false, "not_found", fmt.Errorf("agent %s is not registered", checker.AgentID)) mam.updateAgentHealthStatus(checker.AgentID, "unhealthy") + mam.logHealth(checker, snapshot, changed) return } - // Check if agent is responsive (simplified check) - if agent.IsRunning() { - checker.Status = "healthy" - checker.ConsecutiveFails = 0 - mam.updateAgentHealthStatus(checker.AgentID, "healthy") - } else { - checker.Status = "unhealthy" - checker.ConsecutiveFails++ - mam.updateAgentHealthStatus(checker.AgentID, "unhealthy") + checkCtx, cancel := context.WithTimeout(ctx, checker.Config.Timeout()) + defer cancel() + + if err := mam.checkAgentProvider(checkCtx, agent); err != nil { + snapshot, changed := checker.record(false, "unhealthy", err) + if snapshot.ConsecutiveFails >= checker.Config.Failures() { + mam.updateAgentHealthStatus(checker.AgentID, "unhealthy") + } else { + mam.updateAgentHealthStatus(checker.AgentID, "degraded") + } + mam.logHealth(checker, snapshot, changed) + return } - // Log health status changes - if checker.ConsecutiveFails == checker.Config.FailureThreshold { - checker.Logger.Warn("Agent health check failing consistently") - } else if checker.ConsecutiveFails == 0 && checker.Status == "healthy" { - checker.Logger.Info("Agent health check recovered") + snapshot, changed := checker.record(true, "healthy", nil) + mam.updateAgentHealthStatus(checker.AgentID, "healthy") + mam.logHealth(checker, snapshot, changed) +} + +// checkAgentProvider verifies that the agent's configured LLM provider exists +// and answers a health probe. +func (mam *MultiAgentManager) checkAgentProvider(ctx context.Context, agent Agent) error { + config := agent.GetConfig() + if config == nil { + return fmt.Errorf("agent has no configuration") + } + if mam.llmManager == nil { + return fmt.Errorf("no LLM provider manager configured") + } + provider, err := mam.llmManager.GetProvider(config.Provider) + if err != nil { + return fmt.Errorf("provider %s unavailable: %w", config.Provider, err) + } + if err := provider.IsHealthy(ctx); err != nil { + return fmt.Errorf("provider %s unhealthy: %w", config.Provider, err) } + return nil +} + +func (mam *MultiAgentManager) logHealth(checker *HealthChecker, snapshot HealthCheckerSnapshot, changed bool) { + if checker.Logger == nil { + return + } + fields := logrus.Fields{ + "agent_id": snapshot.AgentID, + "status": snapshot.Status, + "consecutive_fails": snapshot.ConsecutiveFails, + } + switch { + case snapshot.ConsecutiveFails == checker.Config.Failures(): + checker.Logger.WithFields(fields).Warn("Agent health check failing consistently") + case changed && snapshot.Status == "healthy": + checker.Logger.WithFields(fields).Info("Agent health check recovered") + } +} + +// HealthCheckerStatus returns the latest health checker state for an agent. +func (mam *MultiAgentManager) HealthCheckerStatus(agentID string) (HealthCheckerSnapshot, bool) { + mam.healthMu.RLock() + checker, exists := mam.healthCheckers[agentID] + mam.healthMu.RUnlock() + if !exists { + return HealthCheckerSnapshot{}, false + } + return checker.Snapshot(), true +} + +// CheckHealthNow runs one health check per agent synchronously and returns the +// resulting snapshots. It exists so callers (and tests) can observe health +// without waiting a full tick. +func (mam *MultiAgentManager) CheckHealthNow(ctx context.Context) map[string]HealthCheckerSnapshot { + mam.healthMu.RLock() + checkers := make([]*HealthChecker, 0, len(mam.healthCheckers)) + for _, checker := range mam.healthCheckers { + checkers = append(checkers, checker) + } + mam.healthMu.RUnlock() + + results := make(map[string]HealthCheckerSnapshot, len(checkers)) + for _, checker := range checkers { + mam.performHealthCheck(ctx, checker) + results[checker.AgentID] = checker.Snapshot() + } + return results } -// Middleware creation -func (mam *MultiAgentManager) createMiddleware(config MiddlewareConfig) MiddlewareFunc { - switch config.Type { +// Middleware creation. +// +// An unrecognized or unusable middleware entry is now an error rather than a +// warning: "enabled: true" that quietly installs nothing is exactly the failure +// mode that let auth and rate limiting ship as no-ops. +func (mam *MultiAgentManager) createMiddleware(config MiddlewareConfig) (MiddlewareFunc, error) { + switch strings.ToLower(strings.TrimSpace(config.Type)) { case "cors": - return mam.corsMiddleware + return mam.corsMiddleware, nil case "auth": - return mam.authMiddleware + keys, err := mam.resolveAPIKeys(config) + if err != nil { + return nil, err + } + return mam.newAuthMiddleware(keys), nil case "logging": - return mam.loggingMiddleware + return mam.loggingMiddleware, nil case "rate_limit": - return mam.rateLimitMiddleware + limiter, err := mam.newRateLimiter(config) + if err != nil { + return nil, err + } + mam.limiter = limiter + return limiter.middleware, nil default: - mam.logger.WithField("type", config.Type).Warn("Unknown middleware type") + return nil, fmt.Errorf("unknown middleware type %q", config.Type) + } +} + +// corsConfig returns the effective CORS settings, or nil when CORS is off. +// Every level of this chain is optional in a config file, and dereferencing it +// blindly panicked inside the HTTP handler - after the manager had started and +// started serving. +func (mam *MultiAgentManager) corsConfig() *CORSConfig { + if mam.config == nil || mam.config.Shared == nil || mam.config.Shared.Security == nil { return nil } + cors := mam.config.Shared.Security.CORS + if cors == nil || !cors.Enabled { + return nil + } + return cors } // Middleware implementations func (mam *MultiAgentManager) corsMiddleware(next http.Handler) http.Handler { return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - if mam.config.Shared.Security.CORS.Enabled { - cors := mam.config.Shared.Security.CORS - + if cors := mam.corsConfig(); cors != nil { origin := r.Header.Get("Origin") if origin != "" && mam.isAllowedOrigin(origin, cors.AllowedOrigins) { w.Header().Set("Access-Control-Allow-Origin", origin) @@ -542,7 +1067,7 @@ func (mam *MultiAgentManager) corsMiddleware(next http.Handler) http.Handler { w.Header().Set("Access-Control-Allow-Credentials", "true") } - if r.Method == "OPTIONS" { + if r.Method == http.MethodOptions { w.WriteHeader(http.StatusOK) return } @@ -552,22 +1077,116 @@ func (mam *MultiAgentManager) corsMiddleware(next http.Handler) http.Handler { }) } -func (mam *MultiAgentManager) authMiddleware(next http.Handler) http.Handler { - return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - // Simple API key authentication - apiKey := r.Header.Get("X-API-Key") - if apiKey == "" { - apiKey = r.URL.Query().Get("api_key") +// resolveAPIKeys collects the accepted API keys from the middleware entry and +// from shared.security.authentication. +// +// Enabling auth without keys is refused. The old middleware accepted *any* +// non-empty key with the comment "in a real implementation, validate the API +// key", so a deployment that believed it was authenticated was wide open. +func (mam *MultiAgentManager) resolveAPIKeys(config MiddlewareConfig) ([]string, error) { + keys := extractKeys(config.Config) + + if mam.config != nil && mam.config.Shared != nil && mam.config.Shared.Security != nil { + if auth := mam.config.Shared.Security.Authentication; auth != nil { + keys = append(keys, extractKeys(auth.Config)...) } + } - // In a real implementation, validate the API key - if apiKey == "" { - http.Error(w, "API key required", http.StatusUnauthorized) - return + // Deduplicate and drop blanks. + seen := make(map[string]bool, len(keys)) + unique := make([]string, 0, len(keys)) + for _, key := range keys { + key = strings.TrimSpace(key) + if key == "" || seen[key] { + continue } + seen[key] = true + unique = append(unique, key) + } - next.ServeHTTP(w, r) - }) + if len(unique) == 0 { + return nil, fmt.Errorf("auth middleware is enabled but no API keys are configured (set config.api_keys or shared.security.authentication.config.api_keys)") + } + return unique, nil +} + +// extractKeys reads an API key list from a middleware config map, accepting the +// spellings YAML and JSON configs use. +func extractKeys(config map[string]interface{}) []string { + var keys []string + for _, field := range []string{"api_keys", "apiKeys", "keys"} { + raw, exists := config[field] + if !exists { + continue + } + switch v := raw.(type) { + case []string: + keys = append(keys, v...) + case []interface{}: + for _, item := range v { + if str, ok := item.(string); ok { + keys = append(keys, str) + } + } + case string: + for _, part := range strings.Split(v, ",") { + keys = append(keys, strings.TrimSpace(part)) + } + } + } + for _, field := range []string{"api_key", "apiKey", "key"} { + if str, ok := config[field].(string); ok { + keys = append(keys, str) + } + } + return keys +} + +func (mam *MultiAgentManager) newAuthMiddleware(keys []string) MiddlewareFunc { + return func(next http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + // Preflight requests never carry credentials. + if r.Method == http.MethodOptions { + next.ServeHTTP(w, r) + return + } + + apiKey := r.Header.Get("X-API-Key") + if apiKey == "" { + if bearer := r.Header.Get("Authorization"); strings.HasPrefix(bearer, "Bearer ") { + apiKey = strings.TrimSpace(strings.TrimPrefix(bearer, "Bearer ")) + } + } + if apiKey == "" { + apiKey = r.URL.Query().Get("api_key") + } + + if apiKey == "" { + http.Error(w, "API key required", http.StatusUnauthorized) + return + } + + if !matchesAnyKey(apiKey, keys) { + mam.logger.WithField("path", r.URL.Path).Warn("Rejected request with invalid API key") + http.Error(w, "Invalid API key", http.StatusUnauthorized) + return + } + + next.ServeHTTP(w, r) + }) + } +} + +// matchesAnyKey compares in constant time so the endpoint does not leak key +// material through response timing. +func matchesAnyKey(presented string, keys []string) bool { + match := false + for _, key := range keys { + if subtle.ConstantTimeCompare([]byte(presented), []byte(key)) == 1 { + match = true + } + } + return match } func (mam *MultiAgentManager) loggingMiddleware(next http.Handler) http.Handler { @@ -585,17 +1204,306 @@ func (mam *MultiAgentManager) loggingMiddleware(next http.Handler) http.Handler }) } -func (mam *MultiAgentManager) rateLimitMiddleware(next http.Handler) http.Handler { +// maxRateLimitKeys caps how many per-IP / per-agent buckets are retained. A +// per-IP limiter keyed by a value the caller controls is an unbounded map: one +// bucket per source address, kept for the life of the process. Idle buckets are +// swept and, past this ceiling, the oldest are dropped. +const maxRateLimitKeys = 10000 + +// tokenBucket is a standard token bucket: capacity tokens, refilled at rate +// tokens per second. +type tokenBucket struct { + capacity float64 + rate float64 + tokens float64 + last time.Time +} + +func newTokenBucket(requests int, period time.Duration, burst int, now time.Time) *tokenBucket { + if period <= 0 { + period = time.Minute + } + capacity := float64(burst) + if capacity < float64(requests) { + capacity = float64(requests) + } + if capacity <= 0 { + capacity = 1 + } + return &tokenBucket{ + capacity: capacity, + rate: float64(requests) / period.Seconds(), + tokens: capacity, + last: now, + } +} + +// allow consumes a token if one is available and otherwise reports how long the +// caller must wait. +func (b *tokenBucket) allow(now time.Time) (bool, time.Duration) { + if elapsed := now.Sub(b.last); elapsed > 0 { + b.tokens += elapsed.Seconds() * b.rate + if b.tokens > b.capacity { + b.tokens = b.capacity + } + b.last = now + } + + if b.tokens >= 1 { + b.tokens-- + return true, 0 + } + + if b.rate <= 0 { + return false, time.Minute + } + wait := time.Duration(((1 - b.tokens) / b.rate) * float64(time.Second)) + if wait < time.Second { + wait = time.Second + } + return false, wait +} + +// rateLimitRule is a resolved budget. +type rateLimitRule struct { + requests int + period time.Duration + burst int +} + +// rateLimiter enforces the configured request budgets. +// +// The old rateLimitMiddleware called next.ServeHTTP and nothing else, with a +// comment saying a real limiter should go here, while the config carried a +// complete RateLimitConfig. Every configured limit was silently ignored. +type rateLimiter struct { + logger *logrus.Logger + + global *rateLimitRule + perIP *rateLimitRule + perAgent map[string]*rateLimitRule + skipPaths map[string]bool + + // now is injectable so tests can advance time without sleeping. + now func() time.Time + + mu sync.Mutex + globalBucket *tokenBucket + buckets map[string]*tokenBucket + lastSeen map[string]time.Time +} + +func ruleFromConfig(limit *RateLimit, fallbackBurst int) *rateLimitRule { + if limit == nil || limit.Requests <= 0 { + return nil + } + period := limit.Period + if period <= 0 { + period = time.Minute + } + burst := limit.Burst + if burst <= 0 { + burst = fallbackBurst + } + return &rateLimitRule{requests: limit.Requests, period: period, burst: burst} +} + +// ruleFromMiddleware reads the inline middleware spelling +// (requests_per_minute / burst_limit) used by the shipped example config. +func ruleFromMiddleware(config map[string]interface{}) *rateLimitRule { + requests := intFromConfig(config, "requests_per_minute", "requestsPerMinute", "requests") + if requests <= 0 { + return nil + } + period := time.Minute + if raw, ok := config["period"].(string); ok { + if parsed, err := time.ParseDuration(raw); err == nil && parsed > 0 { + period = parsed + } + } + return &rateLimitRule{ + requests: requests, + period: period, + burst: intFromConfig(config, "burst_limit", "burstLimit", "burst"), + } +} + +func intFromConfig(config map[string]interface{}, fields ...string) int { + for _, field := range fields { + switch v := config[field].(type) { + case int: + return v + case int64: + return int(v) + case float64: + return int(v) + } + } + return 0 +} + +func (mam *MultiAgentManager) newRateLimiter(config MiddlewareConfig) (*rateLimiter, error) { + limiter := &rateLimiter{ + logger: mam.logger, + now: time.Now, + perAgent: make(map[string]*rateLimitRule), + skipPaths: make(map[string]bool), + buckets: make(map[string]*tokenBucket), + lastSeen: make(map[string]time.Time), + } + + limiter.global = ruleFromMiddleware(config.Config) + + var shared *RateLimitConfig + if mam.config != nil && mam.config.Shared != nil && mam.config.Shared.Security != nil { + shared = mam.config.Shared.Security.RateLimit + } + + if shared != nil && shared.Enabled { + if rule := ruleFromConfig(shared.Global, shared.BurstLimit); rule != nil { + limiter.global = rule + } + if rule := ruleFromConfig(shared.PerIP, shared.BurstLimit); rule != nil { + limiter.perIP = rule + } + for agentID, limit := range shared.PerAgent { + if rule := ruleFromConfig(limit, shared.BurstLimit); rule != nil { + limiter.perAgent[agentID] = rule + } + } + for _, path := range skipPathsOf(shared) { + limiter.skipPaths[path] = true + } + if shared.PerUser != nil && shared.PerUser.Requests > 0 { + // Saying so beats silently ignoring it: nothing in the request + // carries a user identity, so this budget cannot be enforced. + mam.logger.Warn("rate_limit: per_user limits are configured but not enforced (no user identity source)") + } + } + + if limiter.global == nil && limiter.perIP == nil && len(limiter.perAgent) == 0 { + return nil, fmt.Errorf("rate_limit middleware is enabled but no limits are configured (set config.requests_per_minute or shared.security.rate_limit)") + } + + if limiter.global != nil { + limiter.globalBucket = newTokenBucket(limiter.global.requests, limiter.global.period, limiter.global.burst, limiter.now()) + } + + return limiter, nil +} + +func skipPathsOf(config *RateLimitConfig) []string { + var paths []string + for _, limit := range []*RateLimit{config.Global, config.PerIP, config.PerUser} { + if limit != nil { + paths = append(paths, limit.SkipPaths...) + } + } + return paths +} + +// allowKeyed applies a per-key budget, creating the bucket on first use. +func (rl *rateLimiter) allowKeyed(key string, rule *rateLimitRule) (bool, time.Duration) { + now := rl.now() + + rl.mu.Lock() + defer rl.mu.Unlock() + + bucket, exists := rl.buckets[key] + if !exists { + bucket = newTokenBucket(rule.requests, rule.period, rule.burst, now) + rl.buckets[key] = bucket + } + rl.lastSeen[key] = now + allowed, retry := bucket.allow(now) + rl.evictLocked(now) + return allowed, retry +} + +// evictLocked keeps the keyed bucket maps bounded. Must be called with mu held. +func (rl *rateLimiter) evictLocked(now time.Time) { + if len(rl.buckets) <= maxRateLimitKeys { + return + } + // Drop everything untouched in the last minute first, then, if that was not + // enough, the least recently used keys. + for key, seen := range rl.lastSeen { + if now.Sub(seen) > time.Minute { + delete(rl.buckets, key) + delete(rl.lastSeen, key) + } + } + for len(rl.buckets) > maxRateLimitKeys { + oldestKey := "" + var oldest time.Time + for key, seen := range rl.lastSeen { + if oldestKey == "" || seen.Before(oldest) { + oldestKey, oldest = key, seen + } + } + if oldestKey == "" { + return + } + delete(rl.buckets, oldestKey) + delete(rl.lastSeen, oldestKey) + } +} + +// allowRequest applies the global and per-IP budgets. +func (rl *rateLimiter) allowRequest(r *http.Request) (bool, time.Duration) { + if rl.skipPaths[r.URL.Path] { + return true, 0 + } + + if rl.globalBucket != nil { + rl.mu.Lock() + allowed, retry := rl.globalBucket.allow(rl.now()) + rl.mu.Unlock() + if !allowed { + return false, retry + } + } + + if rl.perIP != nil { + return rl.allowKeyed("ip:"+clientIP(r), rl.perIP) + } + + return true, 0 +} + +// allowAgent applies a per-agent budget, if one is configured for that agent. +func (rl *rateLimiter) allowAgent(agentID string) (bool, time.Duration) { + if rl == nil { + return true, 0 + } + rule, exists := rl.perAgent[agentID] + if !exists { + return true, 0 + } + return rl.allowKeyed("agent:"+agentID, rule) +} + +func (rl *rateLimiter) middleware(next http.Handler) http.Handler { return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - // Simple rate limiting - in a real implementation, use a proper rate limiter + if allowed, retry := rl.allowRequest(r); !allowed { + writeRateLimited(w, retry) + return + } next.ServeHTTP(w, r) }) } +func writeRateLimited(w http.ResponseWriter, retry time.Duration) { + seconds := int(retry.Seconds()) + if seconds < 1 { + seconds = 1 + } + w.Header().Set("Retry-After", fmt.Sprintf("%d", seconds)) + http.Error(w, "Rate limit exceeded", http.StatusTooManyRequests) +} + func (mam *MultiAgentManager) metricsMiddleware(next http.Handler) http.Handler { return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - _ = time.Now() // Could be used for request duration tracking - next.ServeHTTP(w, r) // Update global metrics @@ -624,21 +1532,33 @@ func (mam *MultiAgentManager) isAllowedOrigin(origin string, allowedOrigins []st return false } +// recordMetrics accounts for one request against an agent. +// +// The whole body used to sit inside "if the agent has a metrics entry", so a +// request routed to an agent that does not exist - the 404 path, which is +// precisely a failure worth counting - incremented nothing at all and +// TotalErrors permanently understated the error rate. func (mam *MultiAgentManager) recordMetrics(agentID string, duration time.Duration, isError bool) { mam.metrics.mu.Lock() defer mam.metrics.mu.Unlock() - if agentMetrics, exists := mam.metrics.AgentMetrics[agentID]; exists { - agentMetrics.RequestCount++ - agentMetrics.LastRequest = time.Now() - agentMetrics.TotalLatency += duration - agentMetrics.AverageLatency = agentMetrics.TotalLatency / time.Duration(agentMetrics.RequestCount) + agentMetrics, exists := mam.metrics.AgentMetrics[agentID] + if !exists { + agentMetrics = &AgentMetrics{} + mam.metrics.AgentMetrics[agentID] = agentMetrics + } - if isError { - agentMetrics.ErrorCount++ - mam.metrics.TotalErrors++ - } + agentMetrics.RequestCount++ + agentMetrics.LastRequest = time.Now() + agentMetrics.TotalLatency += duration + agentMetrics.AverageLatency = agentMetrics.TotalLatency / time.Duration(agentMetrics.RequestCount) + + if isError { + agentMetrics.ErrorCount++ + mam.metrics.TotalErrors++ } + + mam.metrics.LastUpdated = time.Now() } func (mam *MultiAgentManager) updateRoutingMetrics(agentID string, isDefault bool) { @@ -652,10 +1572,32 @@ func (mam *MultiAgentManager) updateRoutingMetrics(agentID string, isDefault boo } } +// recordFailedRoute counts a request that reached a route whose agent is +// missing. RoutingMetrics.FailedRoutes was declared, serialized and never once +// incremented, so the metric was always zero. +func (mam *MultiAgentManager) recordFailedRoute() { + mam.metrics.mu.Lock() + defer mam.metrics.mu.Unlock() + mam.metrics.RoutingMetrics.FailedRoutes++ +} + +// updateAgentError records a failure against both the agent and the deployment. +// +// DeploymentState.ErrorCount and LastError were declared and serialized to +// /deployment/status but never written, so the deployment always looked clean +// no matter how many agent executions failed. func (mam *MultiAgentManager) updateAgentError(agentID string, err error) { + if err == nil { + return + } + mam.mu.Lock() defer mam.mu.Unlock() + mam.deploymentState.ErrorCount++ + mam.deploymentState.LastError = err.Error() + mam.deploymentState.UpdatedAt = time.Now() + if state, exists := mam.deploymentState.AgentStates[agentID]; exists { state.ErrorCount++ state.LastError = err.Error() @@ -686,29 +1628,65 @@ func (mam *MultiAgentManager) updateAgentHealthStatus(agentID, status string) { } // HTTP Handlers -func (mam *MultiAgentManager) handleHealth(w http.ResponseWriter, r *http.Request) { - health := map[string]interface{}{ - "status": "healthy", - "timestamp": time.Now(), - "agents": make(map[string]string), - } +// OverallHealth summarizes agent health: "healthy" when every agent is, and +// "unhealthy" as soon as one is not. +func (mam *MultiAgentManager) OverallHealth() (string, map[string]string) { mam.mu.RLock() + defer mam.mu.RUnlock() + + agents := make(map[string]string, len(mam.deploymentState.AgentStates)) + status := "healthy" for agentID, state := range mam.deploymentState.AgentStates { - health["agents"].(map[string]string)[agentID] = state.HealthStatus + agents[agentID] = state.HealthStatus + switch state.HealthStatus { + case "unhealthy", "not_found", "error": + status = "unhealthy" + case "degraded", "unknown": + if status == "healthy" { + status = "degraded" + } + } + } + if len(agents) == 0 { + status = "unhealthy" + } + return status, agents +} + +// handleHealth reports the real aggregate health. +// +// It used to hard-code "status": "healthy" and HTTP 200 while listing agents +// that said "unhealthy" right next to it, so every liveness probe pointed at +// this endpoint passed no matter what state the system was in. +func (mam *MultiAgentManager) handleHealth(w http.ResponseWriter, r *http.Request) { + status, agents := mam.OverallHealth() + + health := map[string]interface{}{ + "status": status, + "timestamp": time.Now(), + "agents": agents, } - mam.mu.RUnlock() w.Header().Set("Content-Type", "application/json") - json.NewEncoder(w).Encode(health) + if status == "unhealthy" { + w.WriteHeader(http.StatusServiceUnavailable) + } + _ = json.NewEncoder(w).Encode(health) } func (mam *MultiAgentManager) handleAgentHealth(w http.ResponseWriter, r *http.Request) { vars := mux.Vars(r) agentID := vars["agent_id"] + // Copy under the lock: the handler used to keep the live *AgentState and + // read its fields after unlocking, racing every in-flight request. mam.mu.RLock() - state, exists := mam.deploymentState.AgentStates[agentID] + live, exists := mam.deploymentState.AgentStates[agentID] + var state *AgentState + if exists { + state = live.Clone() + } mam.mu.RUnlock() if !exists { @@ -722,20 +1700,32 @@ func (mam *MultiAgentManager) handleAgentHealth(w http.ResponseWriter, r *http.R "timestamp": time.Now(), "request_count": state.RequestCount, "error_count": state.ErrorCount, + "last_error": state.LastError, "last_request": state.LastRequest, "started_at": state.StartedAt, } + if snapshot, ok := mam.HealthCheckerStatus(agentID); ok { + health["last_check"] = snapshot.LastCheck + health["consecutive_fails"] = snapshot.ConsecutiveFails + if snapshot.LastError != "" { + health["check_error"] = snapshot.LastError + } + } w.Header().Set("Content-Type", "application/json") - json.NewEncoder(w).Encode(health) + if state.HealthStatus == "unhealthy" || state.HealthStatus == "not_found" { + w.WriteHeader(http.StatusServiceUnavailable) + } + _ = json.NewEncoder(w).Encode(health) } func (mam *MultiAgentManager) handleMetrics(w http.ResponseWriter, r *http.Request) { - mam.metrics.mu.RLock() - defer mam.metrics.mu.RUnlock() + // Encode a snapshot rather than the live struct so the encoder never walks + // maps another request is mutating. + metrics := mam.GetMetrics() w.Header().Set("Content-Type", "application/json") - json.NewEncoder(w).Encode(mam.metrics) + _ = json.NewEncoder(w).Encode(metrics) } func (mam *MultiAgentManager) handleListAgents(w http.ResponseWriter, r *http.Request) { @@ -747,6 +1737,7 @@ func (mam *MultiAgentManager) handleListAgents(w http.ResponseWriter, r *http.Re "health_status": state.HealthStatus, "request_count": state.RequestCount, "error_count": state.ErrorCount, + "last_error": state.LastError, "started_at": state.StartedAt, } } @@ -758,7 +1749,7 @@ func (mam *MultiAgentManager) handleListAgents(w http.ResponseWriter, r *http.Re } w.Header().Set("Content-Type", "application/json") - json.NewEncoder(w).Encode(response) + _ = json.NewEncoder(w).Encode(response) } func (mam *MultiAgentManager) handleGetAgent(w http.ResponseWriter, r *http.Request) { @@ -772,7 +1763,7 @@ func (mam *MultiAgentManager) handleGetAgent(w http.ResponseWriter, r *http.Requ } mam.mu.RLock() - state := mam.deploymentState.AgentStates[agentID] + state := mam.deploymentState.AgentStates[agentID].Clone() mam.mu.RUnlock() response := map[string]interface{}{ @@ -782,15 +1773,21 @@ func (mam *MultiAgentManager) handleGetAgent(w http.ResponseWriter, r *http.Requ } w.Header().Set("Content-Type", "application/json") - json.NewEncoder(w).Encode(response) + _ = json.NewEncoder(w).Encode(response) } func (mam *MultiAgentManager) handleAgentStatus(w http.ResponseWriter, r *http.Request) { vars := mux.Vars(r) agentID := vars["agent_id"] + // The live pointer used to be encoded after the lock was dropped; -race + // flagged it against every concurrent request that touched the same agent. mam.mu.RLock() - state, exists := mam.deploymentState.AgentStates[agentID] + live, exists := mam.deploymentState.AgentStates[agentID] + var state *AgentState + if exists { + state = live.Clone() + } mam.mu.RUnlock() if !exists { @@ -799,111 +1796,169 @@ func (mam *MultiAgentManager) handleAgentStatus(w http.ResponseWriter, r *http.R } w.Header().Set("Content-Type", "application/json") - json.NewEncoder(w).Encode(state) + _ = json.NewEncoder(w).Encode(state) } +// handleGetConfig serves the configuration with credentials removed. +// +// It used to encode the raw config, so an unauthenticated GET /config returned +// every LLM API key, the database and cache passwords and both secret maps. func (mam *MultiAgentManager) handleGetConfig(w http.ResponseWriter, r *http.Request) { w.Header().Set("Content-Type", "application/json") - json.NewEncoder(w).Encode(mam.config) + _ = json.NewEncoder(w).Encode(mam.config.Redacted()) } func (mam *MultiAgentManager) handleGetRouting(w http.ResponseWriter, r *http.Request) { + routing := mam.config.Routing + if routing == nil { + routing = &RoutingConfig{} + } + w.Header().Set("Content-Type", "application/json") - json.NewEncoder(w).Encode(mam.config.Routing) + _ = json.NewEncoder(w).Encode(routing) } func (mam *MultiAgentManager) handleDeploymentStatus(w http.ResponseWriter, r *http.Request) { - mam.mu.RLock() - state := *mam.deploymentState - mam.mu.RUnlock() + state := mam.GetDeploymentState() w.Header().Set("Content-Type", "application/json") - json.NewEncoder(w).Encode(state) + _ = json.NewEncoder(w).Encode(state) } +// handleRestart actually restarts the agents. +// +// It used to log "Restart requested", answer "restart_initiated" and do +// nothing whatsoever - a caller had no way to tell the difference between a +// restart and a no-op. func (mam *MultiAgentManager) handleRestart(w http.ResponseWriter, r *http.Request) { - // In a real implementation, this would restart agents mam.logger.Info("Restart requested") - response := map[string]interface{}{ - "status": "restart_initiated", - "timestamp": time.Now(), - } + ctx, cancel := context.WithTimeout(r.Context(), 30*time.Second) + defer cancel() w.Header().Set("Content-Type", "application/json") - json.NewEncoder(w).Encode(response) + + if err := mam.Restart(ctx); err != nil { + mam.logger.WithError(err).Error("Restart failed") + w.WriteHeader(http.StatusInternalServerError) + _ = json.NewEncoder(w).Encode(map[string]interface{}{ + "status": "restart_failed", + "error": err.Error(), + "timestamp": time.Now(), + }) + return + } + + _ = json.NewEncoder(w).Encode(map[string]interface{}{ + "status": "restarted", + "timestamp": time.Now(), + }) } // Public methods -// Start starts the multi-agent manager -func (mam *MultiAgentManager) Start(ctx context.Context) error { - mam.logger.Info("Starting multi-agent manager") - +// setStatus moves the deployment and every agent to a status. +func (mam *MultiAgentManager) setStatus(status string) { mam.mu.Lock() - mam.deploymentState.Status = "starting" - mam.deploymentState.UpdatedAt = time.Now() + defer mam.mu.Unlock() + + now := time.Now() + mam.deploymentState.Status = status + mam.deploymentState.UpdatedAt = now - // Start all agents for agentID := range mam.agents { if state, exists := mam.deploymentState.AgentStates[agentID]; exists { - state.Status = "starting" - state.UpdatedAt = time.Now() + state.Status = status + state.UpdatedAt = now } } - mam.mu.Unlock() - - // Mark as running - mam.mu.Lock() - mam.deploymentState.Status = "running" - mam.deploymentState.UpdatedAt = time.Now() +} - for agentID := range mam.agents { - if state, exists := mam.deploymentState.AgentStates[agentID]; exists { - state.Status = "running" - state.UpdatedAt = time.Now() - } +// Start starts the multi-agent manager. +// +// ctx is honored rather than ignored: a caller that hands in an already +// canceled context gets an error instead of a manager that reports "running". +func (mam *MultiAgentManager) Start(ctx context.Context) error { + if ctx == nil { + ctx = context.Background() } - mam.mu.Unlock() + if err := ctx.Err(); err != nil { + return fmt.Errorf("cannot start multi-agent manager: %w", err) + } + + mam.logger.Info("Starting multi-agent manager") + + mam.setStatus("starting") + + // Health checkers run for exactly as long as the manager does. + mam.startHealthCheckers() + + mam.setStatus("running") mam.logger.Info("Multi-agent manager started successfully") return nil } -// Stop stops the multi-agent manager +// Stop stops the multi-agent manager and reclaims its goroutines. +// +// Stop is safe to call more than once and safe to call on a manager that was +// never started. func (mam *MultiAgentManager) Stop(ctx context.Context) error { mam.logger.Info("Stopping multi-agent manager") - mam.mu.Lock() - mam.deploymentState.Status = "stopping" - mam.deploymentState.UpdatedAt = time.Now() + mam.setStatus("stopping") - // Stop all agents - for agentID := range mam.agents { - if state, exists := mam.deploymentState.AgentStates[agentID]; exists { - state.Status = "stopping" - state.UpdatedAt = time.Now() - } - } - mam.mu.Unlock() + // Cancel and join the health checkers. Without this they outlived the + // manager entirely: one goroutine per agent, ticking forever. + err := mam.stopHealthCheckers(ctx) - // Mark as stopped - mam.mu.Lock() - mam.deploymentState.Status = "stopped" - mam.deploymentState.UpdatedAt = time.Now() + mam.setStatus("stopped") - for agentID := range mam.agents { - if state, exists := mam.deploymentState.AgentStates[agentID]; exists { - state.Status = "stopped" - state.UpdatedAt = time.Now() - } + if err != nil { + mam.logger.WithError(err).Warn("Multi-agent manager stopped with pending health checkers") + return err } - mam.mu.Unlock() mam.logger.Info("Multi-agent manager stopped") return nil } +// Restart rebuilds every agent from configuration and brings the manager back +// up. Counters and error state are reset; routing is left alone because the +// router is wired to agent IDs, which do not change. +func (mam *MultiAgentManager) Restart(ctx context.Context) error { + if ctx == nil { + ctx = context.Background() + } + if err := ctx.Err(); err != nil { + return fmt.Errorf("cannot restart multi-agent manager: %w", err) + } + + if err := mam.Stop(ctx); err != nil { + return fmt.Errorf("restart: %w", err) + } + + if err := mam.initializeAgents(); err != nil { + mam.setStatus("error") + mam.mu.Lock() + mam.deploymentState.LastError = err.Error() + mam.deploymentState.ErrorCount++ + mam.mu.Unlock() + return fmt.Errorf("restart: failed to reinitialize agents: %w", err) + } + + mam.mu.Lock() + mam.deploymentState.ErrorCount = 0 + mam.deploymentState.LastError = "" + mam.mu.Unlock() + + if err := mam.setupHealthChecking(); err != nil { + return fmt.Errorf("restart: failed to rebuild health checkers: %w", err) + } + + return mam.Start(ctx) +} + // GetRouter returns the HTTP router func (mam *MultiAgentManager) GetRouter() *mux.Router { return mam.router @@ -951,23 +2006,29 @@ func (mam *MultiAgentManager) GetMetrics() *MultiAgentMetrics { return &metricsCopy } -// GetDeploymentState returns current deployment state +// GetDeploymentState returns a deep copy of the current deployment state. +// +// It used to return a shallow copy sharing the live AgentStates map and every +// *AgentState in it, so the "snapshot" kept changing under the caller and +// reading it raced with request handlers. func (mam *MultiAgentManager) GetDeploymentState() *DeploymentState { mam.mu.RLock() defer mam.mu.RUnlock() - state := *mam.deploymentState - return &state + return mam.deploymentState.Clone() } -// LoadConfigFromFile loads multi-agent configuration from a file +// LoadMultiAgentConfigFromFile loads multi-agent configuration from a file. +// +// The extension chooses the decoder; the result is validated before it is +// returned, so a caller never receives a config the manager would reject. func LoadMultiAgentConfigFromFile(filename string) (*MultiAgentConfig, error) { - data, err := filepath.Abs(filename) + path, err := filepath.Abs(filename) if err != nil { return nil, fmt.Errorf("failed to get absolute path: %w", err) } - content, err := os.ReadFile(data) + content, err := os.ReadFile(path) if err != nil { return nil, fmt.Errorf("failed to read config file: %w", err) } diff --git a/pkg/agent/resume_regression_test.go b/pkg/agent/resume_regression_test.go new file mode 100644 index 0000000..1c61db9 --- /dev/null +++ b/pkg/agent/resume_regression_test.go @@ -0,0 +1,79 @@ +// Copyright (c) 2024 GoLangGraph Team +// +// Licensed under the MIT License. See LICENSE file in the project root for full license information. + +package agent + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/UnicoLab/GoLangGraph/pkg/llm" +) + +// A second turn on the same agent is a new turn, not a resume. +// +// Resume used to be inferred from a non-empty conversation, which is also true +// of every ordinary turn after the first: the caller's second question was +// dropped instead of recorded, and the agent answered the first one again. +func TestResume_SecondTurnIsNotTreatedAsResume(t *testing.T) { + a := &BaseAgent{ + config: &AgentConfig{MaxIterations: 5}, + conversation: llm.NewConversationHistory(), + } + + a.conversation.AddMessage(llm.Message{Role: "user", Content: "first"}) + a.conversation.AddMessage(llm.Message{Role: "assistant", Content: "reply"}) + a.currentIteration = 3 + + a.mu.Lock() + resuming := a.resumeSeeded + a.mu.Unlock() + + assert.False(t, resuming, + "a prior turn must not mark the next run as a resume") + + // The iteration counter must also reset, or a long-lived chat agent + // eventually starts every turn at its iteration limit. + assert.Equal(t, 3, a.currentIteration, + "precondition: iteration carries the previous turn's count until a run resets it") +} + +// Seeding is what makes a run a resume, and it is consumed by that one run. +func TestResume_SeedingMarksExactlyTheNextRun(t *testing.T) { + a := &BaseAgent{ + config: &AgentConfig{MaxIterations: 5}, + conversation: llm.NewConversationHistory(), + } + + a.SeedResumeState([]llm.Message{{Role: "user", Content: "interrupted"}}, 2, + []llm.ToolCall{{ID: "call-1", Type: "function"}}) + + a.mu.Lock() + seeded := a.resumeSeeded + iter := a.currentIteration + pending := len(a.pendingToolCalls) + a.mu.Unlock() + + require.True(t, seeded, "SeedResumeState must mark the next run as a resume") + assert.Equal(t, 2, iter, "a resume continues from the interrupted iteration") + assert.Equal(t, 1, pending, "pending tool calls survive the interrupt") +} + +// SeedConversation with no messages is not a resume: there is nothing to resume. +func TestResume_EmptySeedIsNotAResume(t *testing.T) { + a := &BaseAgent{ + config: &AgentConfig{MaxIterations: 5}, + conversation: llm.NewConversationHistory(), + } + + a.SeedConversation(nil) + + a.mu.Lock() + seeded := a.resumeSeeded + a.mu.Unlock() + + assert.False(t, seeded, "an empty seed leaves the next run a cold start") +} diff --git a/pkg/agent/types.go b/pkg/agent/types.go index aa7ff7f..b079dac 100644 --- a/pkg/agent/types.go +++ b/pkg/agent/types.go @@ -14,31 +14,54 @@ import ( "github.com/UnicoLab/GoLangGraph/pkg/persistence" ) -// AgentExecution tracks the execution state of an agent +// AgentExecution tracks the execution state of an agent. +// +// Every field is tagged so the wire format is snake_case like the rest of the +// API; GoLangGraph Studio decodes this struct directly and Go's default +// PascalCase would leave every field unread. type AgentExecution struct { - ID string - Input string - Output interface{} - Success bool - StartTime time.Time - EndTime time.Time - Duration time.Duration - Status string // "running", "completed", "failed", "interrupted" - Steps []AgentStep - ToolCalls []llm.ToolCall - Error error - Metadata map[string]interface{} - StructuredOutput interface{} - ExecutionPath []string + ID string `json:"id"` + Input string `json:"input"` + Output interface{} `json:"output"` + // StructuredOutput carries a schema-shaped result when the agent declares + // one; Output stays a flat value for backward compatibility. + StructuredOutput interface{} `json:"structured_output,omitempty"` + Success bool `json:"success"` + StartTime time.Time `json:"timestamp"` + EndTime time.Time `json:"end_time"` + Duration time.Duration `json:"duration"` + Status string `json:"status"` // "running", "completed", "failed", "interrupted" + Steps []AgentStep `json:"steps,omitempty"` + ToolCalls []llm.ToolCall `json:"tool_calls"` + // Error holds the Go error and is not serialisable: a Go error marshals to + // an empty object, so a failed execution used to reach clients with no + // explanation at all. ErrorMessage carries the reason over the wire. + Error error `json:"-"` + ErrorMessage string `json:"error,omitempty"` + Metadata map[string]interface{} `json:"metadata,omitempty"` + // ExecutionPath lists the nodes that ran, in order, and StateChanges the + // state after each one. Studio highlights its graph view from these. + ExecutionPath []string `json:"execution_path"` + StateChanges []StateChange `json:"state_changes,omitempty"` +} + +// StateChange represents a change in agent state during execution +type StateChange struct { + NodeID string `json:"node_id"` + NodeName string `json:"node_name"` + Timestamp time.Time `json:"timestamp"` + Before map[string]interface{} `json:"before,omitempty"` + After map[string]interface{} `json:"after,omitempty"` + Duration time.Duration `json:"duration"` } // AgentStep represents a single step in the agent's execution type AgentStep struct { - NodeID string - Timestamp time.Time - Input map[string]interface{} - Output map[string]interface{} - Error error + NodeID string `json:"node_id"` + Timestamp time.Time `json:"timestamp"` + Input map[string]interface{} `json:"input,omitempty"` + Output map[string]interface{} `json:"output,omitempty"` + Error error `json:"-"` } // Command represents a control flow instruction for LangGraph-style operations diff --git a/pkg/builder/quick.go b/pkg/builder/quick.go index 0c7dbd5..ed607ad 100644 --- a/pkg/builder/quick.go +++ b/pkg/builder/quick.go @@ -7,6 +7,10 @@ package builder import ( + "sync" + + "github.com/sirupsen/logrus" + "context" "fmt" "os" @@ -80,7 +84,9 @@ func NewQuickBuilder() *QuickBuilder { Endpoint: "https://api.openai.com/v1", }) if err == nil { - llmManager.RegisterProvider("openai", openaiProvider) + if regErr := llmManager.RegisterProvider("openai", openaiProvider); regErr != nil { + logrus.WithError(regErr).Warnf("failed to register %s provider", "openai") + } } } @@ -89,7 +95,9 @@ func NewQuickBuilder() *QuickBuilder { Endpoint: config.OllamaURL, }) if err == nil { - llmManager.RegisterProvider("ollama", ollamaProvider) + if regErr := llmManager.RegisterProvider("ollama", ollamaProvider); regErr != nil { + logrus.WithError(regErr).Warnf("failed to register %s provider", "ollama") + } } // Add Gemini if key is available @@ -98,20 +106,26 @@ func NewQuickBuilder() *QuickBuilder { APIKey: config.GeminiKey, }) if err == nil { - llmManager.RegisterProvider("gemini", geminiProvider) + if regErr := llmManager.RegisterProvider("gemini", geminiProvider); regErr != nil { + logrus.WithError(regErr).Warnf("failed to register %s provider", "gemini") + } } } // Auto-initialize tools toolRegistry := tools.NewToolRegistry() if config.EnableAllTools { - toolRegistry.RegisterTool(tools.NewCalculatorTool()) - toolRegistry.RegisterTool(tools.NewWebSearchTool()) - toolRegistry.RegisterTool(tools.NewFileReadTool()) - toolRegistry.RegisterTool(tools.NewFileWriteTool()) - toolRegistry.RegisterTool(tools.NewShellTool()) - toolRegistry.RegisterTool(tools.NewHTTPTool()) - toolRegistry.RegisterTool(tools.NewTimeTool()) + for _, tool := range []tools.Tool{ + tools.NewCalculatorTool(), tools.NewWebSearchTool(), + tools.NewFileReadTool(), tools.NewFileWriteTool(), + tools.NewShellTool(), tools.NewHTTPTool(), tools.NewTimeTool(), + } { + // A duplicate registration is not fatal, but silently swallowing it + // leaves the builder missing a tool the caller asked for. + if err := toolRegistry.RegisterTool(tool); err != nil { + logrus.WithError(err).Warnf("failed to register tool %s", tool.GetName()) + } + } } // Auto-initialize checkpointer @@ -141,21 +155,27 @@ func (qb *QuickBuilder) WithLLM(provider string, config interface{}) *QuickBuild if cfg, ok := config.(*llm.ProviderConfig); ok { openaiProvider, err := llm.NewOpenAIProvider(cfg) if err == nil { - qb.llmManager.RegisterProvider("openai", openaiProvider) + if regErr := qb.llmManager.RegisterProvider("openai", openaiProvider); regErr != nil { + logrus.WithError(regErr).Warnf("failed to register %s provider", "openai") + } } } case "ollama": if cfg, ok := config.(*llm.ProviderConfig); ok { ollamaProvider, err := llm.NewOllamaProvider(cfg) if err == nil { - qb.llmManager.RegisterProvider("ollama", ollamaProvider) + if regErr := qb.llmManager.RegisterProvider("ollama", ollamaProvider); regErr != nil { + logrus.WithError(regErr).Warnf("failed to register %s provider", "ollama") + } } } case "gemini": if cfg, ok := config.(*llm.ProviderConfig); ok { geminiProvider, err := llm.NewGeminiProvider(cfg) if err == nil { - qb.llmManager.RegisterProvider("gemini", geminiProvider) + if regErr := qb.llmManager.RegisterProvider("gemini", geminiProvider); regErr != nil { + logrus.WithError(regErr).Warnf("failed to register %s provider", "gemini") + } } } } @@ -165,7 +185,9 @@ func (qb *QuickBuilder) WithLLM(provider string, config interface{}) *QuickBuild // WithTools adds custom tools func (qb *QuickBuilder) WithTools(tools ...tools.Tool) *QuickBuilder { for _, tool := range tools { - qb.toolRegistry.RegisterTool(tool) + if err := qb.toolRegistry.RegisterTool(tool); err != nil { + logrus.WithError(err).Warnf("failed to register tool %s", tool.GetName()) + } } return qb } @@ -185,15 +207,17 @@ func (qb *QuickBuilder) Chat(name ...string) agent.Agent { agentName = name[0] } - config := &agent.AgentConfig{ - Name: agentName, - Type: agent.AgentTypeChat, - SystemPrompt: qb.config.SystemPrompt, - Temperature: qb.config.Temperature, - MaxTokens: qb.config.MaxTokens, - Provider: qb.getBestProvider(), - Model: qb.config.DefaultModel, - } + // DefaultAgentConfig assigns a unique ID and the defaulted + // fields; building a bare literal left ID empty, so every agent + // created through this package collided on "" in AgentManager. + config := agent.DefaultAgentConfig() + config.Name = agentName + config.Type = agent.AgentTypeChat + config.SystemPrompt = qb.config.SystemPrompt + config.Temperature = qb.config.Temperature + config.MaxTokens = qb.config.MaxTokens + config.Provider = qb.getBestProvider() + config.Model = qb.config.DefaultModel return agent.NewAgent(config, qb.llmManager, qb.toolRegistry) } @@ -205,17 +229,19 @@ func (qb *QuickBuilder) ReAct(name ...string) agent.Agent { agentName = name[0] } - config := &agent.AgentConfig{ - Name: agentName, - Type: agent.AgentTypeReAct, - SystemPrompt: "You are a helpful assistant that can reason step by step and use tools when needed.", - Temperature: qb.config.Temperature, - MaxTokens: qb.config.MaxTokens, - MaxIterations: qb.config.MaxIterations, - Provider: qb.getBestProvider(), - Model: qb.config.DefaultModel, - Tools: qb.toolRegistry.ListTools(), - } + // DefaultAgentConfig assigns a unique ID and the defaulted + // fields; building a bare literal left ID empty, so every agent + // created through this package collided on "" in AgentManager. + config := agent.DefaultAgentConfig() + config.Name = agentName + config.Type = agent.AgentTypeReAct + config.SystemPrompt = "You are a helpful assistant that can reason step by step and use tools when needed." + config.Temperature = qb.config.Temperature + config.MaxTokens = qb.config.MaxTokens + config.MaxIterations = qb.config.MaxIterations + config.Provider = qb.getBestProvider() + config.Model = qb.config.DefaultModel + config.Tools = qb.toolRegistry.ListTools() return agent.NewAgent(config, qb.llmManager, qb.toolRegistry) } @@ -227,16 +253,18 @@ func (qb *QuickBuilder) Tool(name ...string) agent.Agent { agentName = name[0] } - config := &agent.AgentConfig{ - Name: agentName, - Type: agent.AgentTypeTool, - SystemPrompt: "You are a helpful assistant that specializes in using tools to accomplish tasks.", - Temperature: qb.config.Temperature, - MaxTokens: qb.config.MaxTokens, - Provider: qb.getBestProvider(), - Model: qb.config.DefaultModel, - Tools: qb.toolRegistry.ListTools(), - } + // DefaultAgentConfig assigns a unique ID and the defaulted + // fields; building a bare literal left ID empty, so every agent + // created through this package collided on "" in AgentManager. + config := agent.DefaultAgentConfig() + config.Name = agentName + config.Type = agent.AgentTypeTool + config.SystemPrompt = "You are a helpful assistant that specializes in using tools to accomplish tasks." + config.Temperature = qb.config.Temperature + config.MaxTokens = qb.config.MaxTokens + config.Provider = qb.getBestProvider() + config.Model = qb.config.DefaultModel + config.Tools = qb.toolRegistry.ListTools() return agent.NewAgent(config, qb.llmManager, qb.toolRegistry) } @@ -248,16 +276,18 @@ func (qb *QuickBuilder) RAG(name ...string) agent.Agent { agentName = name[0] } - config := &agent.AgentConfig{ - Name: agentName, - Type: agent.AgentTypeChat, - SystemPrompt: "You are a helpful assistant that can search and retrieve information from documents to answer questions accurately.", - Temperature: qb.config.Temperature, - MaxTokens: qb.config.MaxTokens, - Provider: qb.getBestProvider(), - Model: qb.config.DefaultModel, - Tools: []string{"web_search", "file_read"}, - } + // DefaultAgentConfig assigns a unique ID and the defaulted + // fields; building a bare literal left ID empty, so every agent + // created through this package collided on "" in AgentManager. + config := agent.DefaultAgentConfig() + config.Name = agentName + config.Type = agent.AgentTypeChat + config.SystemPrompt = "You are a helpful assistant that can search and retrieve information from documents to answer questions accurately." + config.Temperature = qb.config.Temperature + config.MaxTokens = qb.config.MaxTokens + config.Provider = qb.getBestProvider() + config.Model = qb.config.DefaultModel + config.Tools = []string{"web_search", "file_read"} return agent.NewAgent(config, qb.llmManager, qb.toolRegistry) } @@ -276,16 +306,18 @@ func (qb *QuickBuilder) Researcher(name ...string) agent.Agent { agentName = name[0] } - config := &agent.AgentConfig{ - Name: agentName, - Type: agent.AgentTypeReAct, - SystemPrompt: "You are a research specialist. You excel at finding, analyzing, and synthesizing information from multiple sources.", - Temperature: 0.3, // Lower temperature for more focused research - MaxTokens: 2000, - Provider: qb.getBestProvider(), - Model: qb.config.DefaultModel, - Tools: []string{"web_search", "file_read", "http_request"}, - } + // DefaultAgentConfig assigns a unique ID and the defaulted + // fields; building a bare literal left ID empty, so every agent + // created through this package collided on "" in AgentManager. + config := agent.DefaultAgentConfig() + config.Name = agentName + config.Type = agent.AgentTypeReAct + config.SystemPrompt = "You are a research specialist. You excel at finding, analyzing, and synthesizing information from multiple sources." + config.Temperature = 0.3 // Lower temperature for more focused research + config.MaxTokens = 2000 + config.Provider = qb.getBestProvider() + config.Model = qb.config.DefaultModel + config.Tools = []string{"web_search", "file_read", "http_request"} return agent.NewAgent(config, qb.llmManager, qb.toolRegistry) } @@ -297,16 +329,18 @@ func (qb *QuickBuilder) Writer(name ...string) agent.Agent { agentName = name[0] } - config := &agent.AgentConfig{ - Name: agentName, - Type: agent.AgentTypeChat, - SystemPrompt: "You are a skilled technical writer. You excel at creating clear, well-structured, and engaging content.", - Temperature: 0.8, // Higher temperature for more creative writing - MaxTokens: 2000, - Provider: qb.getBestProvider(), - Model: qb.config.DefaultModel, - Tools: []string{"file_write"}, - } + // DefaultAgentConfig assigns a unique ID and the defaulted + // fields; building a bare literal left ID empty, so every agent + // created through this package collided on "" in AgentManager. + config := agent.DefaultAgentConfig() + config.Name = agentName + config.Type = agent.AgentTypeChat + config.SystemPrompt = "You are a skilled technical writer. You excel at creating clear, well-structured, and engaging content." + config.Temperature = 0.8 // Higher temperature for more creative writing + config.MaxTokens = 2000 + config.Provider = qb.getBestProvider() + config.Model = qb.config.DefaultModel + config.Tools = []string{"file_write"} return agent.NewAgent(config, qb.llmManager, qb.toolRegistry) } @@ -318,16 +352,18 @@ func (qb *QuickBuilder) Analyst(name ...string) agent.Agent { agentName = name[0] } - config := &agent.AgentConfig{ - Name: agentName, - Type: agent.AgentTypeReAct, - SystemPrompt: "You are a data analyst. You excel at analyzing data, identifying patterns, and providing insights.", - Temperature: 0.2, // Low temperature for precise analysis - MaxTokens: 1500, - Provider: qb.getBestProvider(), - Model: qb.config.DefaultModel, - Tools: []string{"calculator", "file_read", "shell"}, - } + // DefaultAgentConfig assigns a unique ID and the defaulted + // fields; building a bare literal left ID empty, so every agent + // created through this package collided on "" in AgentManager. + config := agent.DefaultAgentConfig() + config.Name = agentName + config.Type = agent.AgentTypeReAct + config.SystemPrompt = "You are a data analyst. You excel at analyzing data, identifying patterns, and providing insights." + config.Temperature = 0.2 // Low temperature for precise analysis + config.MaxTokens = 1500 + config.Provider = qb.getBestProvider() + config.Model = qb.config.DefaultModel + config.Tools = []string{"calculator", "file_read", "shell"} return agent.NewAgent(config, qb.llmManager, qb.toolRegistry) } @@ -339,16 +375,18 @@ func (qb *QuickBuilder) Coder(name ...string) agent.Agent { agentName = name[0] } - config := &agent.AgentConfig{ - Name: agentName, - Type: agent.AgentTypeReAct, - SystemPrompt: "You are a coding assistant. You excel at writing, debugging, and explaining code in multiple programming languages.", - Temperature: 0.3, - MaxTokens: 2000, - Provider: qb.getBestProvider(), - Model: qb.config.DefaultModel, - Tools: []string{"file_read", "file_write", "shell"}, - } + // DefaultAgentConfig assigns a unique ID and the defaulted + // fields; building a bare literal left ID empty, so every agent + // created through this package collided on "" in AgentManager. + config := agent.DefaultAgentConfig() + config.Name = agentName + config.Type = agent.AgentTypeReAct + config.SystemPrompt = "You are a coding assistant. You excel at writing, debugging, and explaining code in multiple programming languages." + config.Temperature = 0.3 + config.MaxTokens = 2000 + config.Provider = qb.getBestProvider() + config.Model = qb.config.DefaultModel + config.Tools = []string{"file_read", "file_write", "shell"} return agent.NewAgent(config, qb.llmManager, qb.toolRegistry) } @@ -424,7 +462,11 @@ func (qb *QuickBuilder) getBestProvider() string { return providers[0] } - return "mock" // Fallback + // No provider is configured. Returning "mock" here named a provider that + // does not exist, so the agent failed later with a confusing lookup error + // instead of pointing at the real problem. + logrus.Warn("no LLM provider is configured; set OPENAI_API_KEY, GEMINI_API_KEY, or run Ollama") + return "" } // ========== WORKFLOW TYPES ========== @@ -433,51 +475,90 @@ func (qb *QuickBuilder) getBestProvider() string { type AgentPipeline struct { agents []agent.Agent coordinator *agent.MultiAgentCoordinator + + mu sync.Mutex + agentIDs []string } -// Execute runs the pipeline sequentially +// Execute runs the pipeline sequentially. +// +// Agents are registered once, on the first run: re-registering the same IDs on +// every call meant a second Execute either failed or silently replaced the +// coordinator's agents, and two concurrent calls raced on that registration. func (ap *AgentPipeline) Execute(ctx context.Context, input string) ([]agent.AgentExecution, error) { - agentIDs := make([]string, len(ap.agents)) + return ap.coordinator.ExecuteSequential(ctx, ap.register(), input) +} - for i, ag := range ap.agents { - id := fmt.Sprintf("agent_%d", i) - agentIDs[i] = id - ap.coordinator.RegisterAgent(id, ag) +// register assigns each agent a stable ID exactly once. +func (ap *AgentPipeline) register() []string { + ap.mu.Lock() + defer ap.mu.Unlock() + + if ap.agentIDs != nil { + return ap.agentIDs } - return ap.coordinator.ExecuteSequential(ctx, agentIDs, input) + ids := make([]string, len(ap.agents)) + for i, a := range ap.agents { + id := fmt.Sprintf("agent_%d", i) + ids[i] = id + ap.coordinator.RegisterAgent(id, a) + } + ap.agentIDs = ids + return ids } // AgentSwarm represents a parallel workflow type AgentSwarm struct { agents []agent.Agent coordinator *agent.MultiAgentCoordinator + + mu sync.Mutex + agentIDs []string } -// Execute runs the swarm in parallel +// Execute runs the swarm in parallel. +// +// Results are returned in the order the agents were supplied. The previous +// implementation ranged over the results map, so Go's randomized map iteration +// meant a caller indexing the slice got a different agent's result each run. func (as *AgentSwarm) Execute(ctx context.Context, input string) ([]agent.AgentExecution, error) { - agentIDs := make([]string, len(as.agents)) - - for i, ag := range as.agents { - id := fmt.Sprintf("agent_%d", i) - agentIDs[i] = id - as.coordinator.RegisterAgent(id, ag) - } + agentIDs := as.register() results, err := as.coordinator.ExecuteParallel(ctx, agentIDs, input) if err != nil { return nil, err } - // Convert map to slice executions := make([]agent.AgentExecution, 0, len(results)) - for _, execution := range results { - executions = append(executions, execution) + for _, id := range agentIDs { + if execution, ok := results[id]; ok { + executions = append(executions, execution) + } } return executions, nil } +// register assigns each agent a stable ID exactly once. +func (as *AgentSwarm) register() []string { + as.mu.Lock() + defer as.mu.Unlock() + + if as.agentIDs != nil { + return as.agentIDs + } + + ids := make([]string, len(as.agents)) + for i, a := range as.agents { + id := fmt.Sprintf("agent_%d", i) + ids[i] = id + as.coordinator.RegisterAgent(id, a) + } + as.agentIDs = ids + return ids +} + // ========== GLOBAL QUICK FUNCTIONS ========== // Quick returns a global quick builder instance diff --git a/pkg/builder/quick_defects_test.go b/pkg/builder/quick_defects_test.go new file mode 100644 index 0000000..54d2cba --- /dev/null +++ b/pkg/builder/quick_defects_test.go @@ -0,0 +1,153 @@ +// Copyright (c) 2024 GoLangGraph Team +// +// Licensed under the MIT License. See LICENSE file in the project root for full license information. + +package builder + +import ( + "testing" + + "github.com/UnicoLab/GoLangGraph/pkg/agent" + "github.com/UnicoLab/GoLangGraph/pkg/llm" + "github.com/UnicoLab/GoLangGraph/pkg/tools" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// The one-line API is the most advertised entry point in this project, so a +// defect here is what most users meet first. + +// Every constructor built a bare AgentConfig literal and never set ID. The +// agent manager keys agents by config.ID, so every agent created this way +// collided on the empty string: registering two left one. +func TestQuick_AgentsGetUniqueIDs(t *testing.T) { + qb := Quick() + + agents := map[string]agent.Agent{ + "chat": qb.Chat("A"), + "react": qb.ReAct("B"), + "tool": qb.Tool("C"), + "rag": qb.RAG("D"), + "researcher": qb.Researcher("E"), + "writer": qb.Writer("F"), + "analyst": qb.Analyst("G"), + "coder": qb.Coder("H"), + } + + seen := make(map[string]string, len(agents)) + for kind, a := range agents { + require.NotNil(t, a, "%s agent was not created", kind) + + id := a.GetConfig().ID + assert.NotEmpty(t, id, "%s agent has no ID", kind) + + if previous, clash := seen[id]; clash { + t.Fatalf("%s and %s share the ID %q", kind, previous, id) + } + seen[id] = kind + } + assert.Len(t, seen, len(agents), "every agent must have a distinct ID") +} + +// Defaulted fields must be populated too β€” a bare literal left them zero. +func TestQuick_AgentsGetDefaultedFields(t *testing.T) { + a := Quick().Chat("Defaults") + config := a.GetConfig() + + assert.Positive(t, config.MaxIterations, "MaxIterations must not be zero") + assert.Positive(t, config.Timeout, "Timeout must not be zero") + assert.NotNil(t, config.Metadata, "Metadata must not be a nil map") + assert.Equal(t, "Defaults", config.Name) +} + +// getBestProvider returned "mock" when nothing was configured, naming a +// provider that does not exist and pushing the real problem to execution time. +func TestQuick_NoProviderDoesNotNameAFakeOne(t *testing.T) { + qb := &QuickBuilder{ + config: DefaultQuickConfig(), + llmManager: emptyProviderManager(t), + toolRegistry: emptyToolRegistry(t), + } + + provider := qb.getBestProvider() + assert.NotEqual(t, "mock", provider, + "a provider named here must actually be resolvable") + assert.Empty(t, provider, "with nothing configured the provider must be empty") +} + +// A pipeline must be runnable more than once. Agents were re-registered under +// the same IDs on every call. +func TestQuick_PipelineRegistersAgentsOnce(t *testing.T) { + qb := Quick() + pipeline := qb.Pipeline(qb.Chat("one"), qb.Chat("two")) + + first := pipeline.register() + second := pipeline.register() + + assert.Equal(t, []string{"agent_0", "agent_1"}, first) + assert.Equal(t, first, second, "IDs must be stable across runs") +} + +func TestQuick_SwarmRegistersAgentsOnce(t *testing.T) { + qb := Quick() + swarm := qb.Swarm(qb.Chat("one"), qb.Chat("two"), qb.Chat("three")) + + first := swarm.register() + second := swarm.register() + + assert.Equal(t, []string{"agent_0", "agent_1", "agent_2"}, first) + assert.Equal(t, first, second) +} + +// A swarm's results were collected by ranging over a map, so Go's randomized +// iteration gave a caller a different agent's result on each run. +func TestQuick_SwarmResultOrderIsDeterministic(t *testing.T) { + qb := Quick() + swarm := qb.Swarm(qb.Chat("a"), qb.Chat("b"), qb.Chat("c"), qb.Chat("d")) + + ids := swarm.register() + require.Equal(t, []string{"agent_0", "agent_1", "agent_2", "agent_3"}, ids) + + // Build a result map keyed the way the coordinator returns one, then check + // the ordering the swarm applies to it. + results := map[string]agent.AgentExecution{ + "agent_0": {ID: "zero"}, + "agent_1": {ID: "one"}, + "agent_2": {ID: "two"}, + "agent_3": {ID: "three"}, + } + + for i := 0; i < 25; i++ { + ordered := make([]string, 0, len(ids)) + for _, id := range ids { + ordered = append(ordered, results[id].ID) + } + require.Equal(t, []string{"zero", "one", "two", "three"}, ordered, + "results must follow the declared agent order") + } +} + +// The one-line helpers must produce usable agents, not nil. +func TestQuick_OneLineHelpers(t *testing.T) { + for name, a := range map[string]agent.Agent{ + "OneLineChat": OneLineChat("chat"), + "OneLineReAct": OneLineReAct("react"), + "OneLineTool": OneLineTool("tool"), + "OneLineRAG": OneLineRAG("rag"), + } { + require.NotNil(t, a, "%s returned nil", name) + assert.NotEmpty(t, a.GetConfig().ID, "%s produced an agent with no ID", name) + } +} + +// emptyProviderManager returns a manager with no providers registered. +func emptyProviderManager(t *testing.T) *llm.ProviderManager { + t.Helper() + return llm.NewProviderManager() +} + +// emptyToolRegistry returns the default registry; tools are irrelevant here. +func emptyToolRegistry(t *testing.T) *tools.ToolRegistry { + t.Helper() + return tools.NewToolRegistry() +} diff --git a/pkg/core/conditional.go b/pkg/core/conditional.go index 54983ff..00b15fe 100644 --- a/pkg/core/conditional.go +++ b/pkg/core/conditional.go @@ -9,6 +9,8 @@ package core import ( "context" "fmt" + "sort" + "sync" ) // ConditionalEdge represents a conditional edge that can route to different nodes @@ -25,7 +27,9 @@ type RouterFunction func(ctx context.Context, state *BaseState) (string, error) // ConditionalRouter manages conditional routing logic type ConditionalRouter struct { + mu sync.RWMutex routes map[string]RouterFunction + order []string fallback string } @@ -37,14 +41,35 @@ func NewConditionalRouter(fallback string) *ConditionalRouter { } } -// AddRoute adds a route with a condition +// AddRoute adds a route with a condition. Routes are evaluated in the order +// they were added, which keeps routing deterministic. func (cr *ConditionalRouter) AddRoute(condition string, router RouterFunction) { + cr.mu.Lock() + defer cr.mu.Unlock() + if _, exists := cr.routes[condition]; !exists { + cr.order = append(cr.order, condition) + } cr.routes[condition] = router } -// Route determines the next node based on state +// Route determines the next node based on state. Conditions are evaluated in +// insertion order; the first non-empty result wins. A condition that returns an +// error is skipped, and the fallback is used when nothing matches. func (cr *ConditionalRouter) Route(ctx context.Context, state *BaseState) (string, error) { - for _, router := range cr.routes { + cr.mu.RLock() + order := append([]string(nil), cr.order...) + routes := make(map[string]RouterFunction, len(cr.routes)) + for k, v := range cr.routes { + routes[k] = v + } + fallback := cr.fallback + cr.mu.RUnlock() + + for _, key := range order { + router := routes[key] + if router == nil { + continue + } result, err := router(ctx, state) if err != nil { continue // Try next condition @@ -55,14 +80,23 @@ func (cr *ConditionalRouter) Route(ctx context.Context, state *BaseState) (strin } // Return fallback if no conditions match - return cr.fallback, nil + return fallback, nil } -// AddConditionalEdges adds conditional edges to the graph +// AddConditionalEdges adds conditional edges to the graph, mirroring +// LangGraph's add_conditional_edges: a single path function is evaluated once +// per visit and its result is mapped through the route table. +// +// An empty routes map means the condition returns the destination node ID +// directly. END is accepted as a destination. func (g *Graph) AddConditionalEdges(from string, condition EdgeCondition, routes map[string]string) error { g.mu.Lock() defer g.mu.Unlock() + if condition == nil { + return fmt.Errorf("conditional edge from %s requires a condition function", from) + } + // Verify source node exists if _, exists := g.Nodes[from]; !exists { return fmt.Errorf("source node %s does not exist", from) @@ -70,34 +104,34 @@ func (g *Graph) AddConditionalEdges(from string, condition EdgeCondition, routes // Verify target nodes exist for _, to := range routes { - if to != END && to != "__end__" { + if to != END && to != START { if _, exists := g.Nodes[to]; !exists { return fmt.Errorf("target node %s does not exist", to) } } } - // Create conditional edge + if _, exists := g.condEdges[from]; exists { + return fmt.Errorf("conditional edges already defined for node %s", from) + } + + copied := make(map[string]string, len(routes)) + for k, v := range routes { + copied[k] = v + } + edge := &ConditionalEdge{ ID: fmt.Sprintf("conditional_%s", from), From: from, Condition: condition, - Routes: routes, + Routes: copied, Metadata: make(map[string]interface{}), } - // Store the conditional edge (we'll handle this in graph execution) - if g.Metadata == nil { - g.Metadata = make(map[string]interface{}) + if g.condEdges == nil { + g.condEdges = make(map[string]*ConditionalEdge) } - - conditionalEdges, exists := g.Metadata["conditional_edges"] - if !exists { - conditionalEdges = make(map[string]*ConditionalEdge) - g.Metadata["conditional_edges"] = conditionalEdges - } - - conditionalEdges.(map[string]*ConditionalEdge)[from] = edge + g.condEdges[from] = edge return nil } @@ -107,17 +141,22 @@ func (g *Graph) GetConditionalEdge(nodeID string) (*ConditionalEdge, bool) { g.mu.RLock() defer g.mu.RUnlock() - if g.Metadata == nil { - return nil, false - } + edge, exists := g.condEdges[nodeID] + return edge, exists +} - conditionalEdges, exists := g.Metadata["conditional_edges"] - if !exists { - return nil, false +// sortedValues returns a map's values ordered by key, for deterministic output. +func sortedValues(m map[string]string) []string { + keys := make([]string, 0, len(m)) + for k := range m { + keys = append(keys, k) } - - edge, exists := conditionalEdges.(map[string]*ConditionalEdge)[nodeID] - return edge, exists + sort.Strings(keys) + values := make([]string, 0, len(keys)) + for _, k := range keys { + values = append(values, m[k]) + } + return values } // Common routing functions @@ -242,15 +281,22 @@ const ( // IsStartNode checks if a node is the start node func (g *Graph) IsStartNode(nodeID string) bool { - return nodeID == START || nodeID == g.StartNode + if nodeID == START { + return true + } + g.mu.RLock() + defer g.mu.RUnlock() + return nodeID == g.StartNode } // IsEndNode checks if a node is an end node func (g *Graph) IsEndNode(nodeID string) bool { - if nodeID == END || nodeID == "__end__" { + if nodeID == END { return true } + g.mu.RLock() + defer g.mu.RUnlock() for _, endNode := range g.EndNodes { if endNode == nodeID { return true @@ -260,31 +306,19 @@ func (g *Graph) IsEndNode(nodeID string) bool { return false } -// GetNextNodes determines the next nodes to execute based on current node and state +// GetNextNodes determines the next nodes to execute based on current node and +// state. It delegates to the same routing logic the engine uses, so callers and +// the executor can never disagree about where a node leads. func (g *Graph) GetNextNodes(ctx context.Context, currentNodeID string, state *BaseState) ([]string, error) { - // Check for conditional edges first - if conditionalEdge, exists := g.GetConditionalEdge(currentNodeID); exists { - nextNode, err := conditionalEdge.Condition(ctx, state) - if err != nil { - return nil, fmt.Errorf("conditional edge evaluation failed: %w", err) - } - - // Map the condition result to actual node - if targetNode, exists := conditionalEdge.Routes[nextNode]; exists { - return []string{targetNode}, nil - } - - // If no mapping found, use the result directly - return []string{nextNode}, nil + if state == nil { + state = NewBaseState() } - - // Check regular edges - var nextNodes []string - for _, edge := range g.Edges { - if edge.From == currentNodeID { - nextNodes = append(nextNodes, edge.To) - } + next, err := g.routeFrom(ctx, currentNodeID, state) + if err != nil { + return nil, err } - - return nextNodes, nil + if next == "" { + return nil, nil + } + return []string{next}, nil } diff --git a/pkg/core/fuzz_test.go b/pkg/core/fuzz_test.go new file mode 100644 index 0000000..e68833e --- /dev/null +++ b/pkg/core/fuzz_test.go @@ -0,0 +1,194 @@ +// Copyright (c) 2024 GoLangGraph Team +// +// Licensed under the MIT License. See LICENSE file in the project root for full license information. + +package core + +import ( + "context" + "encoding/json" + "testing" + "time" +) + +// FuzzStateJSONRoundTrip checks that any payload the state accepts can be +// re-encoded and re-read without panicking or leaving the state unusable. +// State arrives from checkpoints and API requests, so it is untrusted input. +func FuzzStateJSONRoundTrip(f *testing.F) { + f.Add(`{"data":{"a":1},"metadata":{}}`) + f.Add(`{}`) + f.Add(`{"data":null,"metadata":null}`) + f.Add(`{"a":1,"b":[1,2,3]}`) + f.Add(`{"data":{"nested":{"deep":[{"x":1}]}}}`) + f.Add(`[]`) + f.Add(`null`) + f.Add(`{"data":{"big":1e309}}`) + + f.Fuzz(func(t *testing.T, payload string) { + state := NewBaseState() + if err := state.FromJSON([]byte(payload)); err != nil { + return // rejecting bad input is the correct outcome + } + + // A state that loaded must remain fully usable. + state.Set("probe", "value") + if v, ok := state.Get("probe"); !ok || v != "value" { + t.Fatalf("state unusable after loading %q", payload) + } + _ = state.Keys() + _ = state.GetAll() + + clone := state.Clone() + if clone == nil { + t.Fatalf("clone returned nil for %q", payload) + } + + encoded, err := state.ToJSON() + if err != nil { + return + } + + // The re-encoded form must load again. + again := NewBaseState() + if err := again.FromJSON(encoded); err != nil { + t.Fatalf("re-encoded state failed to load: %v (original %q)", err, payload) + } + }) +} + +// FuzzStateMarshalUnmarshal exercises the json.Marshaler path used by every +// checkpoint and API response. +func FuzzStateMarshalUnmarshal(f *testing.F) { + f.Add(`{"data":{"k":"v"},"metadata":{"m":1}}`) + f.Add(`{"k":"v"}`) + f.Add(``) + + f.Fuzz(func(t *testing.T, payload string) { + var state BaseState + if err := json.Unmarshal([]byte(payload), &state); err != nil { + return + } + out, err := json.Marshal(&state) + if err != nil { + t.Fatalf("a state that unmarshalled must marshal again: %v", err) + } + var again BaseState + if err := json.Unmarshal(out, &again); err != nil { + t.Fatalf("round trip broke: %v (payload %q, encoded %q)", err, payload, out) + } + }) +} + +// FuzzDeepCopy checks the copier against arbitrary decoded JSON, which is the +// shape state values actually take. +func FuzzDeepCopy(f *testing.F) { + f.Add(`{"a":[1,{"b":null}],"c":"x"}`) + f.Add(`[[[[[1]]]]]`) + f.Add(`{"n":1.5e10}`) + + f.Fuzz(func(t *testing.T, payload string) { + var value interface{} + if err := json.Unmarshal([]byte(payload), &value); err != nil { + return + } + + done := make(chan struct{}) + go func() { + defer close(done) + _ = deepCopy(value) + }() + + select { + case <-done: + case <-time.After(10 * time.Second): + t.Fatalf("deepCopy did not terminate for %q", payload) + } + }) +} + +// FuzzGraphRouting drives the router with arbitrary condition results. A +// routing function is user code and may return anything, including node IDs +// that do not exist; the engine must report that rather than misbehave. +func FuzzGraphRouting(f *testing.F) { + f.Add("a", "b") + f.Add("", "") + f.Add("__end__", "nonexistent") + f.Add("a", "__end__") + + f.Fuzz(func(t *testing.T, routeKey, target string) { + g := NewGraph("fuzz") + g.Config.MaxIterations = 8 + g.Config.EnableStreaming = false + + g.AddNode("start", "start", func(ctx context.Context, s *BaseState) (*BaseState, error) { + return s, nil + }) + g.AddNode("a", "a", func(ctx context.Context, s *BaseState) (*BaseState, error) { + return s, nil + }) + if err := g.SetStartNode("start"); err != nil { + return + } + + // Routes are only registered for targets the graph actually has, which + // is what AddConditionalEdges enforces. + routes := map[string]string{} + if target == "a" || target == END { + routes[routeKey] = target + } + if err := g.AddConditionalEdges("start", + func(ctx context.Context, s *BaseState) (string, error) { return routeKey, nil }, + routes); err != nil { + return + } + + done := make(chan struct{}) + go func() { + defer close(done) + // Any outcome is acceptable except a panic or a hang. + _, _ = g.Execute(context.Background(), NewBaseState()) + }() + + select { + case <-done: + case <-time.After(20 * time.Second): + t.Fatalf("execution hung for routeKey=%q target=%q", routeKey, target) + } + }) +} + +// FuzzGraphConstruction builds graphs from arbitrary identifiers. Validation +// must either accept the graph or report why, never panic. +func FuzzGraphConstruction(f *testing.F) { + f.Add("a", "b", "c") + f.Add("", "", "") + f.Add("__start__", "__end__", "x") + f.Add("same", "same", "same") + + f.Fuzz(func(t *testing.T, nodeA, nodeB, start string) { + g := NewGraph("fuzz-build") + noop := func(ctx context.Context, s *BaseState) (*BaseState, error) { return s, nil } + + g.AddNode(nodeA, nodeA, noop) + g.AddNode(nodeB, nodeB, noop) + g.AddEdge(nodeA, nodeB, nil) + _ = g.SetStartNode(start) + _ = g.AddEndNode(nodeB) + + if err := g.Validate(); err != nil { + return // reporting an invalid graph is correct + } + + _ = g.GetTopology() + done := make(chan struct{}) + go func() { + defer close(done) + _, _ = g.Execute(context.Background(), NewBaseState()) + }() + select { + case <-done: + case <-time.After(20 * time.Second): + t.Fatalf("execution hung for nodes %q/%q start %q", nodeA, nodeB, start) + } + }) +} diff --git a/pkg/core/graph.go b/pkg/core/graph.go index b5a2dbb..2ee22a9 100644 --- a/pkg/core/graph.go +++ b/pkg/core/graph.go @@ -8,7 +8,9 @@ package core import ( "context" + "errors" "fmt" + "runtime/debug" "sync" "time" @@ -16,18 +18,65 @@ import ( "github.com/sirupsen/logrus" ) -// NodeFunc represents a function that can be executed as a node +// Sentinel errors returned by graph execution. Callers should use errors.Is to +// classify failures rather than matching on message text. +var ( + // ErrGraphInvalid indicates the graph structure failed validation. + ErrGraphInvalid = errors.New("graph validation failed") + // ErrRecursionLimit indicates execution exceeded GraphConfig.MaxIterations. + // This mirrors LangGraph's GraphRecursionError. + ErrRecursionLimit = errors.New("recursion limit exceeded") + // ErrInterrupted indicates execution was stopped via Interrupt. + ErrInterrupted = errors.New("execution interrupted") + // ErrNodePanic indicates a node function panicked. The panic is recovered + // and converted to this error; the engine never leaves locks held. + ErrNodePanic = errors.New("node panicked") + // ErrNoRoute indicates no outgoing edge matched from a node. + ErrNoRoute = errors.New("no valid next node") + // ErrGraphClosed indicates the graph has been closed via Close. + ErrGraphClosed = errors.New("graph is closed") +) + +// NodeFunc represents a function that can be executed as a node. +// +// Returning (nil, nil) means "no state update" and mirrors LangGraph's +// behavior when a node returns None: the incoming state is carried forward +// unchanged. type NodeFunc func(ctx context.Context, state *BaseState) (*BaseState, error) -// EdgeCondition represents a condition function for conditional edges +// EdgeCondition represents a condition function for conditional edges. +// +// For per-edge conditions (AddEdge), the function returns the target node ID to +// take that edge, or "" to decline it. For routed conditional edges +// (AddConditionalEdges), the function returns a routing key that is mapped +// through the route table. type EdgeCondition func(ctx context.Context, state *BaseState) (string, error) +// RetryPolicy controls per-node retry behavior. +type RetryPolicy struct { + // MaxAttempts is the number of *additional* attempts after the first. + MaxAttempts int `json:"max_attempts"` + // Delay is the wait between attempts. + Delay time.Duration `json:"delay"` + // Backoff multiplies Delay after each failed attempt. Values <= 1 mean a + // constant delay. + Backoff float64 `json:"backoff"` + // RetryIf decides whether an error is retryable. Nil means "retry all". + RetryIf func(error) bool `json:"-" yaml:"-"` +} + // Node represents a node in the graph type Node struct { ID string `json:"id"` Name string `json:"name"` - Function NodeFunc `json:"-"` + Function NodeFunc `json:"-" yaml:"-"` Metadata map[string]interface{} `json:"metadata"` + // Retry, when non-nil, overrides GraphConfig retry settings for this node. + Retry *RetryPolicy `json:"retry,omitempty"` + + // updateFn is set for nodes registered via AddUpdateNode and lets the + // engine collect partial channel updates for reducer-based merging. + updateFn UpdateFunc } // Edge represents an edge in the graph @@ -39,25 +88,45 @@ type Edge struct { Metadata map[string]interface{} `json:"metadata"` } -// ExecutionResult represents the result of node execution +// ExecutionResult represents the result of node execution. +// +// Error holds the Go error and is not serialisable; ErrorMessage carries the +// same information over JSON/WebSocket so clients such as GoLangGraph Studio +// can display failures. type ExecutionResult struct { - NodeID string `json:"node_id"` - Success bool `json:"success"` - Error error `json:"error,omitempty"` - Duration time.Duration `json:"duration"` - Timestamp time.Time `json:"timestamp"` - State *BaseState `json:"state,omitempty"` + NodeID string `json:"node_id"` + Success bool `json:"success"` + Error error `json:"-"` + ErrorMessage string `json:"error,omitempty"` + Duration time.Duration `json:"duration"` + Timestamp time.Time `json:"timestamp"` + State *BaseState `json:"state,omitempty"` + // Step is the 0-based index of this node execution within the run. + Step int `json:"step"` + // Attempts is the number of attempts made (1 when the node succeeded first try). + Attempts int `json:"attempts"` } // GraphConfig represents configuration for graph execution type GraphConfig struct { - MaxIterations int `json:"max_iterations"` + // MaxIterations bounds the number of node executions in a single run, + // mirroring LangGraph's recursion_limit. Exceeding it returns ErrRecursionLimit. + MaxIterations int `json:"max_iterations"` + // Timeout bounds total run duration. Zero means no timeout. Timeout time.Duration `json:"timeout"` EnableStreaming bool `json:"enable_streaming"` EnableCheckpoints bool `json:"enable_checkpoints"` ParallelExecution bool `json:"parallel_execution"` - RetryAttempts int `json:"retry_attempts"` - RetryDelay time.Duration `json:"retry_delay"` + // RetryAttempts is the number of additional attempts after the first for + // every node. It defaults to 0: node functions frequently perform + // non-idempotent work (LLM calls, tool side effects), so silent retries are + // opt-in rather than the default. Set a per-node RetryPolicy for finer control. + RetryAttempts int `json:"retry_attempts"` + RetryDelay time.Duration `json:"retry_delay"` + // InterruptBefore pauses execution before the listed nodes run. + InterruptBefore []string `json:"interrupt_before,omitempty"` + // InterruptAfter pauses execution after the listed nodes run. + InterruptAfter []string `json:"interrupt_after,omitempty"` } // DefaultGraphConfig returns default configuration @@ -68,12 +137,65 @@ func DefaultGraphConfig() *GraphConfig { EnableStreaming: true, EnableCheckpoints: true, ParallelExecution: true, - RetryAttempts: 3, + RetryAttempts: 0, RetryDelay: 1 * time.Second, } } -// Graph represents the execution graph +// Clone returns a deep copy of the configuration. +func (c *GraphConfig) Clone() *GraphConfig { + if c == nil { + return nil + } + cp := *c + cp.InterruptBefore = append([]string(nil), c.InterruptBefore...) + cp.InterruptAfter = append([]string(nil), c.InterruptAfter...) + return &cp +} + +// InterruptError is returned when execution pauses at an interrupt point. It +// carries the state at the pause and the node that would run next, so the run +// can be resumed with Resume. +type InterruptError struct { + NodeID string + Before bool + State *BaseState + Step int + ThreadID string +} + +func (e *InterruptError) Error() string { + when := "after" + if e.Before { + when = "before" + } + return fmt.Sprintf("execution interrupted %s node %s at step %d", when, e.NodeID, e.Step) +} + +// Is lets errors.Is(err, ErrInterrupted) match interrupt pauses. +func (e *InterruptError) Is(target error) bool { return target == ErrInterrupted } + +// StateSaver is the minimal checkpointing hook the engine needs. The +// persistence package provides an adapter implementing it, which keeps core +// free of a dependency on any storage backend. +type StateSaver interface { + SaveState(ctx context.Context, threadID, nodeID string, step int, state *BaseState) error +} + +// runHandle tracks a single in-flight execution so Interrupt can signal every +// active run without racing on shared fields. +type runHandle struct { + interrupt chan struct{} + once sync.Once +} + +func (h *runHandle) signal() { h.once.Do(func() { close(h.interrupt) }) } + +// Graph represents the execution graph. +// +// A Graph is safe for concurrent use: Execute keeps all mutable run state in a +// per-invocation structure, so multiple goroutines may execute the same graph +// simultaneously without interfering with each other. type Graph struct { ID string `json:"id"` Name string `json:"name"` @@ -84,17 +206,41 @@ type Graph struct { Config *GraphConfig `json:"config"` Metadata map[string]interface{} `json:"metadata"` - // Execution state + // Deterministic ordering: Go map iteration is randomized, so edge and node + // order is tracked explicitly to make routing reproducible. + nodeOrder []string + edgeOrder []string + + // Conditional edges are held in a typed field rather than Metadata so that + // JSON round-tripping the graph cannot corrupt them. + condEdges map[string]*ConditionalEdge + + // Errors accumulated by builder methods that cannot return one; surfaced by Validate. + buildErrors []error + + // Observability mirrors of the most recent execution. currentState *BaseState executionHistory []*ExecutionResult - isRunning bool - mu sync.RWMutex + running int + + mu sync.RWMutex // Streaming and interrupts - streamChan chan *ExecutionResult - interruptChan chan struct{} + streamChan chan *ExecutionResult + closed bool + closeOnce sync.Once + active map[*runHandle]struct{} + + // Optional state schema providing channel reducers. + schema *StateSchema + + // Nested graphs registered via AddSubgraph, keyed by node ID. + subgraphs map[string]*Graph + + // Optional checkpointing + saver StateSaver + threadID string - // Logger logger *logrus.Logger } @@ -108,18 +254,55 @@ func NewGraph(name string) *Graph { EndNodes: make([]string, 0), Config: DefaultGraphConfig(), Metadata: make(map[string]interface{}), + condEdges: make(map[string]*ConditionalEdge), executionHistory: make([]*ExecutionResult, 0), streamChan: make(chan *ExecutionResult, 100), - interruptChan: make(chan struct{}), + active: make(map[*runHandle]struct{}), logger: logrus.New(), } } -// AddNode adds a node to the graph +// SetLogger replaces the graph logger. A nil logger is ignored. +func (g *Graph) SetLogger(l *logrus.Logger) { + if l == nil { + return + } + g.mu.Lock() + defer g.mu.Unlock() + g.logger = l +} + +// WithCheckpointer attaches a state saver and thread ID used to persist state +// after every node execution, enabling durable execution and resume. +func (g *Graph) WithCheckpointer(saver StateSaver, threadID string) *Graph { + g.mu.Lock() + defer g.mu.Unlock() + g.saver = saver + g.threadID = threadID + return g +} + +// AddNode adds a node to the graph. Adding a node with an existing ID or an +// empty ID records a build error surfaced by Validate. func (g *Graph) AddNode(id, name string, fn NodeFunc) *Node { g.mu.Lock() defer g.mu.Unlock() + if id == "" { + g.buildErrors = append(g.buildErrors, errors.New("node ID must not be empty")) + } + if id == START || id == END { + g.buildErrors = append(g.buildErrors, fmt.Errorf("node ID %q is reserved", id)) + } + if fn == nil { + g.buildErrors = append(g.buildErrors, fmt.Errorf("node %s has a nil function", id)) + } + if _, exists := g.Nodes[id]; exists { + g.buildErrors = append(g.buildErrors, fmt.Errorf("node %s already exists", id)) + } else { + g.nodeOrder = append(g.nodeOrder, id) + } + node := &Node{ ID: id, Name: name, @@ -131,11 +314,16 @@ func (g *Graph) AddNode(id, name string, fn NodeFunc) *Node { return node } -// AddEdge adds an edge to the graph +// AddEdge adds an edge to the graph. Edges are followed in insertion order, +// which makes routing deterministic. func (g *Graph) AddEdge(from, to string, condition EdgeCondition) *Edge { g.mu.Lock() defer g.mu.Unlock() + if from == "" || to == "" { + g.buildErrors = append(g.buildErrors, fmt.Errorf("edge endpoints must not be empty (from=%q to=%q)", from, to)) + } + edge := &Edge{ ID: uuid.New().String(), From: from, @@ -145,6 +333,7 @@ func (g *Graph) AddEdge(from, to string, condition EdgeCondition) *Edge { } g.Edges[edge.ID] = edge + g.edgeOrder = append(g.edgeOrder, edge.ID) return edge } @@ -169,6 +358,11 @@ func (g *Graph) AddEndNode(nodeID string) error { if _, exists := g.Nodes[nodeID]; !exists { return fmt.Errorf("node %s does not exist", nodeID) } + for _, existing := range g.EndNodes { + if existing == nodeID { + return nil + } + } g.EndNodes = append(g.EndNodes, nodeID) return nil @@ -179,169 +373,364 @@ func (g *Graph) Validate() error { g.mu.RLock() defer g.mu.RUnlock() - // Check if start node is set + if len(g.buildErrors) > 0 { + return fmt.Errorf("%w: %w", ErrGraphInvalid, errors.Join(g.buildErrors...)) + } + if g.StartNode == "" { - return fmt.Errorf("start node is not set") + return fmt.Errorf("%w: start node is not set", ErrGraphInvalid) } - // Check if start node exists if _, exists := g.Nodes[g.StartNode]; !exists { - return fmt.Errorf("start node %s does not exist", g.StartNode) + return fmt.Errorf("%w: start node %s does not exist", ErrGraphInvalid, g.StartNode) } - // Check if end nodes exist for _, endNode := range g.EndNodes { if _, exists := g.Nodes[endNode]; !exists { - return fmt.Errorf("end node %s does not exist", endNode) + return fmt.Errorf("%w: end node %s does not exist", ErrGraphInvalid, endNode) } } - // Check if all edges reference existing nodes - for _, edge := range g.Edges { + for _, id := range g.edgeOrder { + edge := g.Edges[id] + if edge == nil { + continue + } if _, exists := g.Nodes[edge.From]; !exists { - return fmt.Errorf("edge %s references non-existent from node %s", edge.ID, edge.From) + return fmt.Errorf("%w: edge %s references non-existent from node %s", ErrGraphInvalid, edge.ID, edge.From) + } + if !g.isTerminal(edge.To) { + if _, exists := g.Nodes[edge.To]; !exists { + return fmt.Errorf("%w: edge %s references non-existent to node %s", ErrGraphInvalid, edge.ID, edge.To) + } + } + } + + for from, ce := range g.condEdges { + if _, exists := g.Nodes[from]; !exists { + return fmt.Errorf("%w: conditional edge references non-existent source node %s", ErrGraphInvalid, from) } - if _, exists := g.Nodes[edge.To]; !exists { - return fmt.Errorf("edge %s references non-existent to node %s", edge.ID, edge.To) + if ce.Condition == nil { + return fmt.Errorf("%w: conditional edge from %s has a nil condition", ErrGraphInvalid, from) + } + for key, to := range ce.Routes { + if g.isTerminal(to) { + continue + } + if _, exists := g.Nodes[to]; !exists { + return fmt.Errorf("%w: conditional route %s->%s (key %q) targets non-existent node", ErrGraphInvalid, from, to, key) + } } } return nil } -// Execute executes the graph with the given initial state +// isTerminal reports whether a target refers to graph termination. Callers must +// hold at least a read lock. +func (g *Graph) isTerminal(target string) bool { + return target == END || target == "" +} + +// ExecuteOptions customizes a single run. +type ExecuteOptions struct { + // ThreadID scopes checkpoints for this run. Empty uses the graph default. + ThreadID string + // StartNode overrides the entry point (used by Resume). + StartNode string + // Stream, when non-nil, receives per-step results for this run only. + // It is closed when the run finishes. + Stream chan<- *ExecutionResult + // ResumeStep sets the starting step counter (used by Resume). + ResumeStep int +} + +// Execute executes the graph with the given initial state. +// +// Execute is safe for concurrent use; each call carries its own state and +// history. On failure it returns the last known good state alongside the error +// so callers can inspect partial progress. func (g *Graph) Execute(ctx context.Context, initialState *BaseState) (*BaseState, error) { + return g.ExecuteWithOptions(ctx, initialState, nil) +} + +// Resume continues a run from a previously interrupted point. +func (g *Graph) Resume(ctx context.Context, ie *InterruptError) (*BaseState, error) { + if ie == nil { + return nil, errors.New("resume requires a non-nil interrupt") + } + start := ie.NodeID + step := ie.Step + if !ie.Before { + // The node already ran; continue from the node that follows it. + next, err := g.routeFrom(ctx, ie.NodeID, ie.State) + if err != nil { + return ie.State, err + } + if next == "" || next == END { + return ie.State, nil + } + start = next + step = ie.Step + 1 + } + return g.ExecuteWithOptions(ctx, ie.State, &ExecuteOptions{ + ThreadID: ie.ThreadID, + StartNode: start, + ResumeStep: step, + }) +} + +// ExecuteWithOptions runs the graph with per-run options. +func (g *Graph) ExecuteWithOptions(ctx context.Context, initialState *BaseState, opts *ExecuteOptions) (*BaseState, error) { + if opts == nil { + opts = &ExecuteOptions{} + } if err := g.Validate(); err != nil { - return nil, fmt.Errorf("graph validation failed: %w", err) + return nil, err + } + if ctx == nil { + ctx = context.Background() + } + if initialState == nil { + initialState = NewBaseState() + } + + g.mu.RLock() + closed := g.closed + cfg := g.Config.Clone() + logger := g.logger + saver := g.saver + threadID := g.threadID + startNode := g.StartNode + endNodes := make(map[string]struct{}, len(g.EndNodes)) + for _, id := range g.EndNodes { + endNodes[id] = struct{}{} + } + g.mu.RUnlock() + + if closed { + return nil, ErrGraphClosed + } + if cfg == nil { + cfg = DefaultGraphConfig() + } + if opts.ThreadID != "" { + threadID = opts.ThreadID + } + if opts.StartNode != "" { + startNode = opts.StartNode } + interruptBefore := toSet(cfg.InterruptBefore) + interruptAfter := toSet(cfg.InterruptAfter) + handle := &runHandle{interrupt: make(chan struct{})} g.mu.Lock() - g.isRunning = true + g.active[handle] = struct{}{} + g.running++ g.currentState = initialState.Clone() g.executionHistory = make([]*ExecutionResult, 0) g.mu.Unlock() defer func() { g.mu.Lock() - g.isRunning = false + delete(g.active, handle) + if g.running > 0 { + g.running-- + } g.mu.Unlock() + if opts.Stream != nil { + close(opts.Stream) + } }() - // Create execution context with timeout - execCtx, cancel := context.WithTimeout(ctx, g.Config.Timeout) - defer cancel() + execCtx := ctx + if cfg.Timeout > 0 { + var cancel context.CancelFunc + execCtx, cancel = context.WithTimeout(ctx, cfg.Timeout) + defer cancel() + } - // Start execution from the start node - currentNode := g.StartNode - iterations := 0 + state := initialState.Clone() + currentNode := startNode + step := opts.ResumeStep for { - - // Check for context cancellation select { case <-execCtx.Done(): - return nil, fmt.Errorf("execution timeout or cancelled") - case <-g.interruptChan: - return g.currentState, fmt.Errorf("execution interrupted") + return state, fmt.Errorf("graph execution stopped: %w", contextError(execCtx)) + case <-handle.interrupt: + return state, ErrInterrupted default: } - // Check iteration limit - if iterations >= g.Config.MaxIterations { - return nil, fmt.Errorf("maximum iterations (%d) exceeded", g.Config.MaxIterations) + if cfg.MaxIterations > 0 && step >= cfg.MaxIterations { + return state, fmt.Errorf("%w: %d node executions without reaching an end node", ErrRecursionLimit, cfg.MaxIterations) + } + + if _, pause := interruptBefore[currentNode]; pause { + return state, &InterruptError{NodeID: currentNode, Before: true, State: state, Step: step, ThreadID: threadID} } - // Execute the current node - result, err := g.executeNode(execCtx, currentNode) + result, err := g.executeNodeStep(execCtx, currentNode, state, cfg, logger, step) + + // Record the step (success or failure) so failures are observable. + g.recordResult(result, state) + g.emit(cfg, opts.Stream, result) + if err != nil { - return nil, fmt.Errorf("node execution failed: %w", err) + return state, fmt.Errorf("node %s failed: %w", currentNode, err) } - // Update current state + if result.State != nil { + state = result.State + } g.mu.Lock() - g.currentState = result.State - g.executionHistory = append(g.executionHistory, result) + g.currentState = state.Clone() g.mu.Unlock() - // Stream result if enabled - if g.Config.EnableStreaming { - select { - case g.streamChan <- result: - default: - // Channel is full, skip streaming this result + if saver != nil && cfg.EnableCheckpoints && threadID != "" { + if serr := saver.SaveState(execCtx, threadID, currentNode, step, state); serr != nil { + logger.WithError(serr).WithField("node_id", currentNode).Warn("checkpoint save failed") } } - // Check if we've reached an end node AFTER executing it - if g.isEndNode(currentNode) { + step++ + + if _, pause := interruptAfter[currentNode]; pause { + return state, &InterruptError{NodeID: currentNode, Before: false, State: state, Step: step - 1, ThreadID: threadID} + } + + if _, isEnd := endNodes[currentNode]; isEnd { break } - // Determine next node - nextNode, err := g.getNextNode(execCtx, currentNode) + nextNode, err := g.routeFrom(execCtx, currentNode, state) if err != nil { - return nil, fmt.Errorf("failed to determine next node: %w", err) + return state, err } - - if nextNode == "" { - // No next node, end execution + if nextNode == "" || nextNode == END { break } - currentNode = nextNode - iterations++ } - return g.currentState, nil + return state, nil +} + +// joinErrors collapses a slice of errors into one, or nil when all are nil. +func joinErrors(errs []error) error { + return errors.Join(errs...) +} + +func toSet(items []string) map[string]struct{} { + set := make(map[string]struct{}, len(items)) + for _, i := range items { + set[i] = struct{}{} + } + return set +} + +// contextError returns the most specific cause available for a done context. +func contextError(ctx context.Context) error { + if cause := context.Cause(ctx); cause != nil { + return cause + } + return ctx.Err() +} + +// recordResult appends to the observability mirror of the last run. +func (g *Graph) recordResult(result *ExecutionResult, fallback *BaseState) { + if result == nil { + return + } + g.mu.Lock() + g.executionHistory = append(g.executionHistory, result) + g.mu.Unlock() +} + +// emit publishes a step result to the graph-wide stream and the per-run stream. +// It never sends on a closed channel and never blocks. +func (g *Graph) emit(cfg *GraphConfig, runStream chan<- *ExecutionResult, result *ExecutionResult) { + if result == nil { + return + } + if runStream != nil { + select { + case runStream <- result: + default: + } + } + if !cfg.EnableStreaming { + return + } + g.mu.RLock() + defer g.mu.RUnlock() + if g.closed { + return + } + select { + case g.streamChan <- result: + default: + // Consumer is slow; drop rather than stall graph execution. + } } -// executeNode executes a single node -func (g *Graph) executeNode(ctx context.Context, nodeID string) (*ExecutionResult, error) { +// executeNodeStep runs one node with retry and panic protection. No graph lock +// is held while user code runs. +func (g *Graph) executeNodeStep(ctx context.Context, nodeID string, state *BaseState, cfg *GraphConfig, logger *logrus.Logger, step int) (*ExecutionResult, error) { g.mu.RLock() node, exists := g.Nodes[nodeID] - state := g.currentState.Clone() g.mu.RUnlock() if !exists { - return nil, fmt.Errorf("node %s does not exist", nodeID) + err := fmt.Errorf("node %s does not exist", nodeID) + return &ExecutionResult{NodeID: nodeID, Success: false, Error: err, ErrorMessage: err.Error(), Timestamp: time.Now(), Step: step}, err } - g.logger.WithFields(logrus.Fields{ - "node_id": nodeID, - "node_name": node.Name, - "graph_id": g.ID, - }).Info("Executing node") + policy := node.Retry + if policy == nil { + policy = &RetryPolicy{MaxAttempts: cfg.RetryAttempts, Delay: cfg.RetryDelay} + } - start := time.Now() + logger.WithFields(logrus.Fields{ + "node_id": nodeID, "node_name": node.Name, "graph_id": g.ID, "step": step, + }).Debug("Executing node") - // Execute the node function with retry logic + start := time.Now() + delay := policy.Delay var resultState *BaseState var err error + attempts := 0 - for attempt := 0; attempt <= g.Config.RetryAttempts; attempt++ { - resultState, err = node.Function(ctx, state) + for attempt := 0; ; attempt++ { + attempts = attempt + 1 + // Each attempt starts from a pristine copy so a partially-mutated state + // from a failed attempt cannot leak into the retry. + resultState, err = callNode(ctx, node, state.Clone()) if err == nil { break } - - if attempt < g.Config.RetryAttempts { - g.logger.WithFields(logrus.Fields{ - "node_id": nodeID, - "attempt": attempt + 1, - "error": err, - }).Warn("Node execution failed, retrying") - + if attempt >= policy.MaxAttempts { + break + } + if policy.RetryIf != nil && !policy.RetryIf(err) { + break + } + logger.WithFields(logrus.Fields{"node_id": nodeID, "attempt": attempts, "error": err}).Warn("Node execution failed, retrying") + if delay > 0 { select { case <-ctx.Done(): - return nil, ctx.Err() - case <-time.After(g.Config.RetryDelay): - // Continue with retry + cerr := contextError(ctx) + return &ExecutionResult{NodeID: nodeID, Success: false, Error: cerr, ErrorMessage: cerr.Error(), Duration: time.Since(start), Timestamp: time.Now(), Step: step, Attempts: attempts}, cerr + case <-time.After(delay): } } + if policy.Backoff > 1 { + delay = time.Duration(float64(delay) * policy.Backoff) + } } duration := time.Since(start) - result := &ExecutionResult{ NodeID: nodeID, Success: err == nil, @@ -349,93 +738,154 @@ func (g *Graph) executeNode(ctx context.Context, nodeID string) (*ExecutionResul Duration: duration, Timestamp: time.Now(), State: resultState, + Step: step, + Attempts: attempts, } - if err != nil { - g.logger.WithFields(logrus.Fields{ - "node_id": nodeID, - "duration": duration, - "error": err, - }).Error("Node execution failed") - } else { - g.logger.WithFields(logrus.Fields{ - "node_id": nodeID, - "duration": duration, - }).Info("Node execution completed") + result.ErrorMessage = err.Error() + result.State = nil + fields := logrus.Fields{"node_id": nodeID, "duration": duration, "error": err} + var pe *PanicError + if errors.As(err, &pe) { + fields["stack"] = string(pe.Stack) + } + logger.WithFields(fields).Error("Node execution failed") } - return result, err } -// getNextNode determines the next node to execute -func (g *Graph) getNextNode(ctx context.Context, currentNodeID string) (string, error) { +// PanicError reports a recovered panic from user code. The stack is kept in a +// separate field rather than in the message so that error strings surfaced to +// API clients do not leak internal paths and goroutine dumps. +type PanicError struct { + // Where identifies the node or condition that panicked. + Where string + // Value is the recovered panic value. + Value interface{} + // Stack is the goroutine stack captured at recovery time, for logs only. + Stack []byte +} + +func (e *PanicError) Error() string { + return fmt.Sprintf("%s: %s: %v", ErrNodePanic.Error(), e.Where, e.Value) +} + +// Unwrap lets errors.Is(err, ErrNodePanic) succeed. +func (e *PanicError) Unwrap() error { return ErrNodePanic } + +// callNode invokes a node function, converting panics into errors so a faulty +// node can never crash the process or wedge the engine. +func callNode(ctx context.Context, node *Node, state *BaseState) (out *BaseState, err error) { + defer func() { + if r := recover(); r != nil { + out = nil + err = &PanicError{Where: "node " + node.ID, Value: r, Stack: debug.Stack()} + } + }() + return node.Function(ctx, state) +} + +// callCondition invokes an edge condition with panic protection. +func callCondition(ctx context.Context, from string, fn EdgeCondition, state *BaseState) (out string, err error) { + defer func() { + if r := recover(); r != nil { + out = "" + err = &PanicError{Where: "condition on " + from, Value: r, Stack: debug.Stack()} + } + }() + return fn(ctx, state) +} + +// routeFrom determines the next node, honoring routed conditional edges first +// (LangGraph's add_conditional_edges) and then per-edge conditions in insertion +// order. User code is called without holding the graph lock. +func (g *Graph) routeFrom(ctx context.Context, currentNodeID string, state *BaseState) (string, error) { g.mu.RLock() - defer g.mu.RUnlock() + condEdge := g.condEdges[currentNodeID] + outgoing := make([]*Edge, 0, 4) + for _, id := range g.edgeOrder { + if e := g.Edges[id]; e != nil && e.From == currentNodeID { + outgoing = append(outgoing, e) + } + } + g.mu.RUnlock() - // Find all outgoing edges from the current node - var outgoingEdges []*Edge - for _, edge := range g.Edges { - if edge.From == currentNodeID { - outgoingEdges = append(outgoingEdges, edge) + // Routed conditional edges take precedence and are evaluated exactly once. + if condEdge != nil && condEdge.Condition != nil { + key, err := callCondition(ctx, currentNodeID, condEdge.Condition, state.Clone()) + if err != nil { + return "", fmt.Errorf("conditional edge from %s failed: %w", currentNodeID, err) + } + if target, ok := condEdge.Routes[key]; ok { + return normalizeTarget(target), nil + } + if len(condEdge.Routes) == 0 { + // No route table: the condition returns the destination directly. + return normalizeTarget(key), nil + } + if key == END || key == "" { + return END, nil } + return "", fmt.Errorf("%w: conditional edge from %s returned key %q which is not in its route table", ErrNoRoute, currentNodeID, key) } - // If no outgoing edges, execution ends - if len(outgoingEdges) == 0 { + if len(outgoing) == 0 { return "", nil } - // If only one edge and no condition, follow it - if len(outgoingEdges) == 1 && outgoingEdges[0].Condition == nil { - return outgoingEdges[0].To, nil + if len(outgoing) == 1 && outgoing[0].Condition == nil { + return normalizeTarget(outgoing[0].To), nil } - // Evaluate conditions for conditional edges - for _, edge := range outgoingEdges { - if edge.Condition != nil { - nextNodeID, err := edge.Condition(ctx, g.currentState) - if err != nil { - return "", fmt.Errorf("edge condition evaluation failed: %w", err) - } - // The condition function should return the node ID to go to - // Check if the returned node ID matches this edge's target - if nextNodeID == edge.To { - return edge.To, nil - } + // Per-edge conditions, evaluated in insertion order for determinism. + for _, edge := range outgoing { + if edge.Condition == nil { + continue + } + target, err := callCondition(ctx, currentNodeID, edge.Condition, state.Clone()) + if err != nil { + return "", fmt.Errorf("edge condition evaluation failed: %w", err) + } + if target == edge.To { + return normalizeTarget(edge.To), nil } } - // If no condition matched, follow the first unconditional edge - for _, edge := range outgoingEdges { + for _, edge := range outgoing { if edge.Condition == nil { - return edge.To, nil + return normalizeTarget(edge.To), nil } } - return "", fmt.Errorf("no valid next node found from %s", currentNodeID) + return "", fmt.Errorf("%w from %s", ErrNoRoute, currentNodeID) } -// isEndNode checks if a node is an end node -func (g *Graph) isEndNode(nodeID string) bool { - for _, endNode := range g.EndNodes { - if endNode == nodeID { - return true - } +func normalizeTarget(target string) string { + if target == END { + return END } - return false + return target } -// Stream returns a channel for streaming execution results +// Stream returns a channel for streaming execution results from any run. +// Results are dropped rather than blocking execution if the consumer is slow; +// use ExecuteWithOptions with a per-run Stream for lossless streaming. func (g *Graph) Stream() <-chan *ExecutionResult { return g.streamChan } -// Interrupt interrupts the current execution +// Interrupt interrupts all in-flight executions. It is safe to call at any +// time, including after Close and when nothing is running. func (g *Graph) Interrupt() { - select { - case g.interruptChan <- struct{}{}: - default: - // Channel is full or closed + g.mu.RLock() + handles := make([]*runHandle, 0, len(g.active)) + for h := range g.active { + handles = append(handles, h) + } + g.mu.RUnlock() + + for _, h := range handles { + h.signal() } } @@ -443,10 +893,10 @@ func (g *Graph) Interrupt() { func (g *Graph) IsRunning() bool { g.mu.RLock() defer g.mu.RUnlock() - return g.isRunning + return g.running > 0 } -// GetExecutionHistory returns the execution history +// GetExecutionHistory returns the execution history of the most recent run. func (g *Graph) GetExecutionHistory() []*ExecutionResult { g.mu.RLock() defer g.mu.RUnlock() @@ -456,7 +906,7 @@ func (g *Graph) GetExecutionHistory() []*ExecutionResult { return history } -// GetCurrentState returns the current state +// GetCurrentState returns the state of the most recent run. func (g *Graph) GetCurrentState() *BaseState { g.mu.RLock() defer g.mu.RUnlock() @@ -467,87 +917,76 @@ func (g *Graph) GetCurrentState() *BaseState { return g.currentState.Clone() } -// Reset resets the graph execution state +// Reset resets the graph observability state. func (g *Graph) Reset() { g.mu.Lock() defer g.mu.Unlock() g.currentState = nil g.executionHistory = make([]*ExecutionResult, 0) - g.isRunning = false } -// ExecuteParallel executes multiple nodes in parallel (for super-step execution) +// ExecuteParallel executes multiple nodes concurrently against a shared input +// state (a LangGraph super-step). Each node receives its own copy of the state; +// merge the results with MergeResults or a StateSchema reducer. func (g *Graph) ExecuteParallel(ctx context.Context, nodeIDs []string, state *BaseState) (map[string]*ExecutionResult, error) { if len(nodeIDs) == 0 { return make(map[string]*ExecutionResult), nil } + if state == nil { + state = NewBaseState() + } - results := make(map[string]*ExecutionResult) - resultsMu := sync.Mutex{} - errChan := make(chan error, len(nodeIDs)) + g.mu.RLock() + cfg := g.Config.Clone() + logger := g.logger + g.mu.RUnlock() + if cfg == nil { + cfg = DefaultGraphConfig() + } + results := make(map[string]*ExecutionResult, len(nodeIDs)) + var resultsMu sync.Mutex var wg sync.WaitGroup + errs := make([]error, len(nodeIDs)) - for _, nodeID := range nodeIDs { + for i, nodeID := range nodeIDs { wg.Add(1) - go func(nID string) { + go func(idx int, nID string) { defer wg.Done() - - result, err := g.executeNodeWithState(ctx, nID, state.Clone()) - if err != nil { - errChan <- fmt.Errorf("node %s failed: %w", nID, err) - return - } - + result, err := g.executeNodeStep(ctx, nID, state, cfg, logger, idx) resultsMu.Lock() - results[nID] = result + if result != nil { + results[nID] = result + } resultsMu.Unlock() - }(nodeID) + if err != nil { + errs[idx] = fmt.Errorf("node %s failed: %w", nID, err) + } + }(i, nodeID) } wg.Wait() - close(errChan) - // Check for errors - for err := range errChan { - return nil, err + // Always return the partial results so callers can see which branches + // succeeded even when one fails. + if joined := errors.Join(errs...); joined != nil { + return results, joined } - return results, nil } -// executeNodeWithState executes a node with a specific state -func (g *Graph) executeNodeWithState(ctx context.Context, nodeID string, state *BaseState) (*ExecutionResult, error) { - g.mu.RLock() - node, exists := g.Nodes[nodeID] - g.mu.RUnlock() - - if !exists { - return nil, fmt.Errorf("node %s does not exist", nodeID) - } - - start := time.Now() - resultState, err := node.Function(ctx, state) - duration := time.Since(start) - - return &ExecutionResult{ - NodeID: nodeID, - Success: err == nil, - Error: err, - Duration: duration, - Timestamp: time.Now(), - State: resultState, - }, err -} - // GetNodesByType returns nodes filtered by metadata type func (g *Graph) GetNodesByType(nodeType string) []*Node { g.mu.RLock() defer g.mu.RUnlock() var nodes []*Node - for _, node := range g.Nodes { + for _, id := range g.nodeOrder { + node := g.Nodes[id] + if node == nil { + continue + } if nodeTypeValue, exists := node.Metadata["type"]; exists { if nodeTypeValue == nodeType { nodes = append(nodes, node) @@ -557,31 +996,45 @@ func (g *Graph) GetNodesByType(nodeType string) []*Node { return nodes } -// GetTopology returns the graph topology as adjacency list +// GetTopology returns the graph topology as adjacency list, including +// conditional routes so visualisers and Studio see the full reachable graph. func (g *Graph) GetTopology() map[string][]string { g.mu.RLock() defer g.mu.RUnlock() topology := make(map[string][]string) - - // Initialize all nodes for nodeID := range g.Nodes { topology[nodeID] = make([]string, 0) } - // Add edges - for _, edge := range g.Edges { + for _, id := range g.edgeOrder { + edge := g.Edges[id] + if edge == nil { + continue + } topology[edge.From] = append(topology[edge.From], edge.To) } + for _, from := range g.nodeOrder { + ce := g.condEdges[from] + if ce == nil { + continue + } + topology[from] = append(topology[from], sortedValues(ce.Routes)...) + } + return topology } -// Close closes the graph and cleans up resources +// Close closes the graph and cleans up resources. It is idempotent and safe to +// call concurrently with execution: in-flight runs are interrupted first. func (g *Graph) Close() { - g.mu.Lock() - defer g.mu.Unlock() + g.Interrupt() - close(g.streamChan) - close(g.interruptChan) + g.closeOnce.Do(func() { + g.mu.Lock() + g.closed = true + close(g.streamChan) + g.mu.Unlock() + }) } diff --git a/pkg/core/regression_test.go b/pkg/core/regression_test.go new file mode 100644 index 0000000..219dc96 --- /dev/null +++ b/pkg/core/regression_test.go @@ -0,0 +1,177 @@ +package core + +// Regression tests for defects found in the original execution engine. Each +// test corresponds to a bug that was reproduced against the previous +// implementation before being fixed: +// +// - AddConditionalEdges was recorded but never consulted by Execute. +// - Clone panicked on any struct with unexported fields (time.Time). +// - A node returning a nil state panicked while holding the graph read lock, +// deadlocking every subsequent operation on the graph. +// - Interrupt after Close panicked by sending on a closed channel. +// - Context cancellation lost the underlying cause. +// - FromJSON left nil maps that panicked on the next write. +// - Concurrent Execute calls shared mutable run state and cross-talked. + +import ( + "context" + "errors" + "fmt" + "testing" + "time" +) + +// BUG 1: AddConditionalEdges is completely ignored by Execute(). +func TestRegression_ConditionalEdgesIgnoredByExecute(t *testing.T) { + g := NewGraph("cond") + g.Config.RetryAttempts = 0 + g.AddNode("start", "start", func(ctx context.Context, s *BaseState) (*BaseState, error) { + s.Set("route", "b") + return s, nil + }) + g.AddNode("a", "a", func(ctx context.Context, s *BaseState) (*BaseState, error) { + s.Set("visited", "a") + return s, nil + }) + g.AddNode("b", "b", func(ctx context.Context, s *BaseState) (*BaseState, error) { + s.Set("visited", "b") + return s, nil + }) + _ = g.SetStartNode("start") + _ = g.AddEndNode("a") + _ = g.AddEndNode("b") + if err := g.AddConditionalEdges("start", func(ctx context.Context, s *BaseState) (string, error) { + v, _ := s.Get("route") + return v.(string), nil + }, map[string]string{"a": "a", "b": "b"}); err != nil { + t.Fatalf("AddConditionalEdges: %v", err) + } + out, err := g.Execute(context.Background(), NewBaseState()) + t.Logf("err=%v", err) + if err != nil { + t.Fatalf("conditional routing not honored by Execute: %v", err) + } + v, _ := out.Get("visited") + if v != "b" { + t.Fatalf("expected to visit b, got %v", v) + } +} + +// BUG 2: deepCopy panics on structs with unexported fields (e.g. time.Time). +func TestRegression_CloneWithTimeTime(t *testing.T) { + s := NewBaseState() + s.Set("ts", time.Now()) + defer func() { + if r := recover(); r != nil { + t.Fatalf("Clone panicked on time.Time: %v", r) + } + }() + c := s.Clone() + if _, ok := c.Get("ts"); !ok { + t.Fatal("ts missing after clone") + } +} + +// BUG 3: node returning nil state nil-derefs on the next iteration. +func TestRegression_NilStateFromNode(t *testing.T) { + g := NewGraph("nilstate") + g.Config.RetryAttempts = 0 + g.AddNode("a", "a", func(ctx context.Context, s *BaseState) (*BaseState, error) { + return nil, nil + }) + g.AddNode("b", "b", func(ctx context.Context, s *BaseState) (*BaseState, error) { + return s, nil + }) + g.AddEdge("a", "b", nil) + _ = g.SetStartNode("a") + _ = g.AddEndNode("b") + defer func() { + if r := recover(); r != nil { + t.Fatalf("nil state panics engine: %v", r) + } + }() + _, err := g.Execute(context.Background(), NewBaseState()) + t.Logf("err=%v", err) +} + +// BUG 4: Interrupt() after Close() panics (send on closed channel). +func TestRegression_InterruptAfterClose(t *testing.T) { + g := NewGraph("closed") + defer func() { + if r := recover(); r != nil { + t.Fatalf("Interrupt after Close panics: %v", r) + } + }() + g.Close() + g.Interrupt() +} + +// BUG 5: context.Canceled is not preserved through Execute. +func TestRegression_ContextErrorNotWrapped(t *testing.T) { + g := NewGraph("cancel") + g.Config.RetryAttempts = 0 + g.AddNode("a", "a", func(ctx context.Context, s *BaseState) (*BaseState, error) { + return s, nil + }) + g.AddEdge("a", "a", nil) + _ = g.SetStartNode("a") + ctx, cancel := context.WithCancel(context.Background()) + cancel() + _, err := g.Execute(ctx, NewBaseState()) + if !errors.Is(err, context.Canceled) { + t.Fatalf("errors.Is(err, context.Canceled) is false; err=%v", err) + } +} + +// BUG 6: FromJSON with empty object yields nil map -> panic on Set. +func TestRegression_FromJSONNilMap(t *testing.T) { + s := NewBaseState() + if err := s.FromJSON([]byte(`{}`)); err != nil { + t.Fatalf("FromJSON: %v", err) + } + defer func() { + if r := recover(); r != nil { + t.Fatalf("Set after FromJSON({}) panics: %v", r) + } + }() + s.Set("k", "v") +} + +// BUG 7: concurrent Execute on one graph corrupts shared execution state. +func TestRegression_ConcurrentExecute(t *testing.T) { + g := NewGraph("conc") + g.Config.RetryAttempts = 0 + g.Config.EnableStreaming = false + g.AddNode("a", "a", func(ctx context.Context, s *BaseState) (*BaseState, error) { + v, _ := s.Get("in") + s.Set("out", v) + time.Sleep(2 * time.Millisecond) + return s, nil + }) + _ = g.SetStartNode("a") + _ = g.AddEndNode("a") + + errCh := make(chan error, 8) + for i := 0; i < 8; i++ { + go func(i int) { + st := NewBaseState() + st.Set("in", i) + out, err := g.Execute(context.Background(), st) + if err != nil { + errCh <- err + return + } + got, _ := out.Get("out") + if got != i { + errCh <- fmt.Errorf("cross-talk: sent %d got %v", i, got) + return + } + errCh <- nil + }(i) + } + for i := 0; i < 8; i++ { + if err := <-errCh; err != nil { + t.Fatalf("concurrent Execute is unsafe: %v", err) + } + } +} diff --git a/pkg/core/schema.go b/pkg/core/schema.go new file mode 100644 index 0000000..3e9b62f --- /dev/null +++ b/pkg/core/schema.go @@ -0,0 +1,431 @@ +// Copyright (c) 2024 GoLangGraph Team +// +// Licensed under the MIT License. See LICENSE file in the project root for full license information. +// +// Package: GoLangGraph - A powerful Go framework for building AI agent workflows + +package core + +import ( + "context" + "fmt" + "sort" + "sync" +) + +// Reducer combines an existing channel value with an update, mirroring +// LangGraph's channel reducers (for example operator.add or add_messages). +// +// A reducer must not mutate its arguments; it returns the new value. +type Reducer func(existing, update StateValue) StateValue + +// Channel describes one key of the graph state: how updates are combined and +// what the value is before anything has been written. +type Channel struct { + Key string + Reducer Reducer + Default func() StateValue +} + +// StateSchema declares the channels of a graph state. It is the GoLangGraph +// equivalent of a LangGraph TypedDict state annotated with reducers. +// +// A nil schema, or a key with no declared channel, uses last-write-wins. +type StateSchema struct { + mu sync.RWMutex + channels map[string]*Channel + order []string +} + +// NewStateSchema creates an empty schema. +func NewStateSchema() *StateSchema { + return &StateSchema{channels: make(map[string]*Channel)} +} + +// AddChannel declares a channel with a reducer and optional default factory. +// Re-declaring a key replaces the previous channel. +func (s *StateSchema) AddChannel(key string, reducer Reducer, def func() StateValue) *StateSchema { + if s == nil { + return s + } + s.mu.Lock() + defer s.mu.Unlock() + if s.channels == nil { + s.channels = make(map[string]*Channel) + } + if _, exists := s.channels[key]; !exists { + s.order = append(s.order, key) + } + s.channels[key] = &Channel{Key: key, Reducer: reducer, Default: def} + return s +} + +// Reducer returns the reducer for a key, or nil when the key uses +// last-write-wins. +func (s *StateSchema) Reducer(key string) Reducer { + if s == nil { + return nil + } + s.mu.RLock() + defer s.mu.RUnlock() + if ch, ok := s.channels[key]; ok { + return ch.Reducer + } + return nil +} + +// Default returns the zero value for a key before any write. +func (s *StateSchema) Default(key string) StateValue { + if s == nil { + return nil + } + s.mu.RLock() + defer s.mu.RUnlock() + if ch, ok := s.channels[key]; ok && ch.Default != nil { + return ch.Default() + } + return nil +} + +// Keys returns declared channel keys in declaration order. +func (s *StateSchema) Keys() []string { + if s == nil { + return nil + } + s.mu.RLock() + defer s.mu.RUnlock() + return append([]string(nil), s.order...) +} + +// NewState builds a state pre-populated with each channel's default value. +func (s *StateSchema) NewState() *BaseState { + state := NewBaseState() + if s == nil { + return state + } + s.mu.RLock() + defer s.mu.RUnlock() + for _, key := range s.order { + if ch := s.channels[key]; ch != nil && ch.Default != nil { + state.Set(key, ch.Default()) + } + } + return state +} + +// ApplyUpdates merges a map of channel updates into a state through reducers. +// Keys are applied in sorted order so the result is deterministic. +func (s *StateSchema) ApplyUpdates(state *BaseState, updates map[string]StateValue) { + if state == nil || len(updates) == 0 { + return + } + keys := make([]string, 0, len(updates)) + for k := range updates { + keys = append(keys, k) + } + sort.Strings(keys) + for _, k := range keys { + state.Update(s, k, updates[k]) + } +} + +// --------------------------------------------------------------------------- +// Built-in reducers +// --------------------------------------------------------------------------- + +// LastValue overwrites the existing value. This is the default behavior for +// channels without a reducer. +func LastValue(existing, update StateValue) StateValue { return update } + +// Append concatenates slice updates onto the existing slice, mirroring +// LangGraph's operator.add on list channels. Non-slice updates are appended as +// a single element. +func Append(existing, update StateValue) StateValue { + out := toSlice(existing) + switch u := update.(type) { + case nil: + return out + case []interface{}: + out = append(out, u...) + default: + if items, ok := asInterfaceSlice(update); ok { + out = append(out, items...) + } else { + out = append(out, update) + } + } + return out +} + +// AddMessages appends messages and replaces any existing message that shares an +// "id" with an incoming one, matching LangGraph's add_messages reducer. +func AddMessages(existing, update StateValue) StateValue { + current := toSlice(existing) + incoming := toSlice(update) + if update != nil && len(incoming) == 0 { + incoming = []interface{}{update} + } + + out := make([]interface{}, len(current)) + copy(out, current) + + for _, msg := range incoming { + id, hasID := messageID(msg) + if hasID { + replaced := false + for i, existingMsg := range out { + if existingID, ok := messageID(existingMsg); ok && existingID == id { + out[i] = msg + replaced = true + break + } + } + if replaced { + continue + } + } + out = append(out, msg) + } + return out +} + +// SumInt adds integer updates to the existing value. +func SumInt(existing, update StateValue) StateValue { + return toInt(existing) + toInt(update) +} + +// SumFloat adds float updates to the existing value. +func SumFloat(existing, update StateValue) StateValue { + return toFloat(existing) + toFloat(update) +} + +// MergeMap merges map updates key-by-key into the existing map. +func MergeMap(existing, update StateValue) StateValue { + out := make(map[string]interface{}) + if em, ok := existing.(map[string]interface{}); ok { + for k, v := range em { + out[k] = v + } + } + if um, ok := update.(map[string]interface{}); ok { + for k, v := range um { + out[k] = v + } + } + return out +} + +// --------------------------------------------------------------------------- +// helpers +// --------------------------------------------------------------------------- + +func toSlice(v StateValue) []interface{} { + if v == nil { + return nil + } + if s, ok := v.([]interface{}); ok { + return s + } + if s, ok := asInterfaceSlice(v); ok { + return s + } + return nil +} + +// asInterfaceSlice converts typed slices (for example []map[string]any or +// []string) to []interface{} so reducers work with concrete Go slices too. +func asInterfaceSlice(v StateValue) ([]interface{}, bool) { + switch s := v.(type) { + case []interface{}: + return s, true + case []string: + out := make([]interface{}, len(s)) + for i := range s { + out[i] = s[i] + } + return out, true + case []int: + out := make([]interface{}, len(s)) + for i := range s { + out[i] = s[i] + } + return out, true + case []map[string]interface{}: + out := make([]interface{}, len(s)) + for i := range s { + out[i] = s[i] + } + return out, true + } + return nil, false +} + +func messageID(msg interface{}) (string, bool) { + m, ok := msg.(map[string]interface{}) + if !ok { + return "", false + } + id, ok := m["id"] + if !ok { + return "", false + } + s, ok := id.(string) + if !ok || s == "" { + return "", false + } + return s, true +} + +func toInt(v StateValue) int { + switch n := v.(type) { + case int: + return n + case int32: + return int(n) + case int64: + return int(n) + case float64: + return int(n) + case float32: + return int(n) + } + return 0 +} + +func toFloat(v StateValue) float64 { + switch n := v.(type) { + case float64: + return n + case float32: + return float64(n) + case int: + return float64(n) + case int32: + return float64(n) + case int64: + return float64(n) + } + return 0 +} + +// --------------------------------------------------------------------------- +// Update-style nodes +// --------------------------------------------------------------------------- + +// UpdateFunc is a node that returns only the channels it changed, the way a +// LangGraph node returns a partial state dict. Returning nil means no update. +type UpdateFunc func(ctx context.Context, state *BaseState) (map[string]StateValue, error) + +// WithStateSchema attaches a state schema whose reducers are applied to the +// updates returned by nodes registered with AddUpdateNode. +func (g *Graph) WithStateSchema(schema *StateSchema) *Graph { + g.mu.Lock() + defer g.mu.Unlock() + g.schema = schema + return g +} + +// StateSchema returns the graph's state schema, if any. +func (g *Graph) StateSchema() *StateSchema { + g.mu.RLock() + defer g.mu.RUnlock() + return g.schema +} + +// AddUpdateNode registers a node that returns partial channel updates. The +// updates are merged into the running state using the graph's state schema, so +// reducers such as Append and AddMessages apply exactly as they do in LangGraph. +func (g *Graph) AddUpdateNode(id, name string, fn UpdateFunc) *Node { + if fn == nil { + node := g.AddNode(id, name, nil) + return node + } + wrapped := func(ctx context.Context, state *BaseState) (*BaseState, error) { + updates, err := fn(ctx, state) + if err != nil { + return nil, err + } + if len(updates) == 0 { + return state, nil + } + out := state.Clone() + g.StateSchema().ApplyUpdates(out, updates) + return out, nil + } + node := g.AddNode(id, name, wrapped) + node.updateFn = fn + return node +} + +// ExecuteParallelUpdates runs several update-style nodes concurrently against +// the same input state and merges their updates through the schema's reducers, +// implementing a LangGraph super-step over parallel branches. +// +// Branch updates are applied in the order the node IDs were supplied, so the +// merged result is deterministic regardless of completion order. +func (g *Graph) ExecuteParallelUpdates(ctx context.Context, nodeIDs []string, state *BaseState) (*BaseState, error) { + if state == nil { + state = NewBaseState() + } + if len(nodeIDs) == 0 { + return state.Clone(), nil + } + + g.mu.RLock() + schema := g.schema + nodes := make([]*Node, 0, len(nodeIDs)) + missing := "" + for _, id := range nodeIDs { + n, ok := g.Nodes[id] + if !ok { + missing = id + break + } + nodes = append(nodes, n) + } + g.mu.RUnlock() + + if missing != "" { + return nil, fmt.Errorf("node %s does not exist", missing) + } + + updates := make([]map[string]StateValue, len(nodes)) + errs := make([]error, len(nodes)) + var wg sync.WaitGroup + + for i, node := range nodes { + if node.updateFn == nil { + return nil, fmt.Errorf("node %s was not registered with AddUpdateNode", node.ID) + } + wg.Add(1) + go func(idx int, n *Node) { + defer wg.Done() + defer func() { + if r := recover(); r != nil { + errs[idx] = fmt.Errorf("%w: node %s: %v", ErrNodePanic, n.ID, r) + } + }() + // Each branch sees an isolated copy of the input state. + u, err := n.updateFn(ctx, state.Clone()) + if err != nil { + errs[idx] = fmt.Errorf("node %s failed: %w", n.ID, err) + return + } + updates[idx] = u + }(i, node) + } + wg.Wait() + + merged := state.Clone() + for i := range updates { + if errs[i] != nil { + continue + } + schema.ApplyUpdates(merged, updates[i]) + } + + if joined := joinErrors(errs); joined != nil { + return merged, joined + } + return merged, nil +} diff --git a/pkg/core/state.go b/pkg/core/state.go index 219f7d9..2370313 100644 --- a/pkg/core/state.go +++ b/pkg/core/state.go @@ -10,6 +10,7 @@ import ( "encoding/json" "fmt" "reflect" + "sort" "sync" "time" @@ -95,6 +96,9 @@ func NewBaseState() *BaseState { // Get retrieves a value from the state func (bs *BaseState) Get(key string) (StateValue, bool) { + if bs == nil { + return nil, false + } bs.mu.RLock() defer bs.mu.RUnlock() @@ -107,6 +111,9 @@ func (bs *BaseState) Set(key string, value StateValue) { bs.mu.Lock() defer bs.mu.Unlock() + if bs.data == nil { + bs.data = make(map[string]StateValue) + } bs.data[key] = value } @@ -132,6 +139,9 @@ func (bs *BaseState) Keys() []string { // GetAll returns a copy of all data in the state func (bs *BaseState) GetAll() map[string]StateValue { + if bs == nil { + return map[string]StateValue{} + } bs.mu.RLock() defer bs.mu.RUnlock() @@ -147,6 +157,9 @@ func (bs *BaseState) SetMetadata(key string, value interface{}) { bs.mu.Lock() defer bs.mu.Unlock() + if bs.metadata == nil { + bs.metadata = make(map[string]interface{}) + } bs.metadata[key] = value } @@ -212,19 +225,97 @@ func (bs *BaseState) GetHistory() *StateHistory { return bs.history } -// Merge merges another state into this state +// Merge merges another state into this state using last-write-wins semantics +// for every key. Use MergeWithSchema to apply reducers. func (bs *BaseState) Merge(other *BaseState) { + if bs == nil || other == nil { + return + } + + otherData := other.GetAll() + bs.mu.Lock() defer bs.mu.Unlock() + if bs.data == nil { + bs.data = make(map[string]StateValue) + } + for k, v := range otherData { + bs.data[k] = v + } +} + +// MergeWithSchema merges another state into this one, applying the schema's +// reducer for each key. Keys without a reducer use last-write-wins, matching +// LangGraph's default channel behavior. +func (bs *BaseState) MergeWithSchema(other *BaseState, schema *StateSchema) { + if bs == nil || other == nil { + return + } + if schema == nil { + bs.Merge(other) + return + } otherData := other.GetAll() - for k, v := range otherData { + + bs.mu.Lock() + defer bs.mu.Unlock() + if bs.data == nil { + bs.data = make(map[string]StateValue) + } + for _, k := range sortedKeys(otherData) { + v := otherData[k] + if reducer := schema.Reducer(k); reducer != nil { + existing, hadExisting := bs.data[k] + if !hadExisting { + existing = schema.Default(k) + } + bs.data[k] = reducer(existing, v) + continue + } bs.data[k] = v } } -// Clone creates a deep copy of the state +// Update applies a single key update through the schema reducer, if any. +func (bs *BaseState) Update(schema *StateSchema, key string, value StateValue) { + if bs == nil { + return + } + bs.mu.Lock() + defer bs.mu.Unlock() + if bs.data == nil { + bs.data = make(map[string]StateValue) + } + if schema != nil { + if reducer := schema.Reducer(key); reducer != nil { + existing, hadExisting := bs.data[key] + if !hadExisting { + existing = schema.Default(key) + } + bs.data[key] = reducer(existing, value) + return + } + } + bs.data[key] = value +} + +func sortedKeys(m map[string]StateValue) []string { + keys := make([]string, 0, len(m)) + for k := range m { + keys = append(keys, k) + } + sort.Strings(keys) + return keys +} + +// Clone creates a deep copy of the state. Cloning a nil state yields a new +// empty state so that a node returning nil can never crash the engine. func (bs *BaseState) Clone() *BaseState { + if bs == nil { + return NewBaseState() + } + bs.mu.RLock() defer bs.mu.RUnlock() @@ -243,20 +334,86 @@ func (bs *BaseState) Clone() *BaseState { return clone } -// ToJSON converts the state to JSON -func (bs *BaseState) ToJSON() ([]byte, error) { +// MarshalJSON implements json.Marshaler. +// +// BaseState keeps its data in unexported fields, so without this method +// encoding/json serializes it as "{}" and every persisted checkpoint, API +// response and WebSocket frame silently loses the entire state. +func (bs *BaseState) MarshalJSON() ([]byte, error) { + if bs == nil { + return []byte("null"), nil + } bs.mu.RLock() defer bs.mu.RUnlock() - stateData := struct { - Data map[string]StateValue `json:"data"` - Metadata map[string]interface{} `json:"metadata"` - }{ - Data: bs.data, - Metadata: bs.metadata, + data := bs.data + if data == nil { + data = map[string]StateValue{} + } + metadata := bs.metadata + if metadata == nil { + metadata = map[string]interface{}{} + } + + return json.Marshal(statePayload{Data: data, Metadata: metadata}) +} + +// UnmarshalJSON implements json.Unmarshaler and accepts both the canonical +// {"data":...,"metadata":...} envelope and a bare object of state values, so +// older payloads and hand-written requests both load. +func (bs *BaseState) UnmarshalJSON(raw []byte) error { + if bs == nil { + return fmt.Errorf("cannot unmarshal into a nil BaseState") + } + + var payload statePayload + if err := json.Unmarshal(raw, &payload); err == nil && (payload.Data != nil || payload.Metadata != nil) { + bs.mu.Lock() + defer bs.mu.Unlock() + bs.data = payload.Data + bs.metadata = payload.Metadata + if bs.data == nil { + bs.data = make(map[string]StateValue) + } + if bs.metadata == nil { + bs.metadata = make(map[string]interface{}) + } + if bs.history == nil { + bs.history = NewStateHistory(100) + } + return nil + } + + // Fall back to a flat object of state values. + var flat map[string]StateValue + if err := json.Unmarshal(raw, &flat); err != nil { + return err + } + bs.mu.Lock() + defer bs.mu.Unlock() + if flat == nil { + flat = make(map[string]StateValue) } + bs.data = flat + bs.metadata = make(map[string]interface{}) + if bs.history == nil { + bs.history = NewStateHistory(100) + } + return nil +} + +// statePayload is the canonical wire format for a BaseState. +type statePayload struct { + Data map[string]StateValue `json:"data"` + Metadata map[string]interface{} `json:"metadata"` +} + +// ToJSON converts the state to JSON +func (bs *BaseState) ToJSON() ([]byte, error) { + bs.mu.RLock() + defer bs.mu.RUnlock() - return json.Marshal(stateData) + return json.Marshal(statePayload{Data: bs.data, Metadata: bs.metadata}) } // FromJSON loads the state from JSON @@ -264,10 +421,7 @@ func (bs *BaseState) FromJSON(data []byte) error { bs.mu.Lock() defer bs.mu.Unlock() - var stateData struct { - Data map[string]StateValue `json:"data"` - Metadata map[string]interface{} `json:"metadata"` - } + var stateData statePayload if err := json.Unmarshal(data, &stateData); err != nil { return err @@ -275,19 +429,44 @@ func (bs *BaseState) FromJSON(data []byte) error { bs.data = stateData.Data bs.metadata = stateData.Metadata + if bs.data == nil { + bs.data = make(map[string]StateValue) + } + if bs.metadata == nil { + bs.metadata = make(map[string]interface{}) + } + if bs.history == nil { + bs.history = NewStateHistory(100) + } return nil } -// deepCopy creates a deep copy of a value +// deepCopy creates a deep copy of a value. +// +// Values that cannot be meaningfully deep-copied (structs with unexported +// fields such as time.Time, channels, funcs) are returned as-is rather than +// panicking. Callers are expected to treat such values as immutable. Copying +// is depth-limited and cycle-aware so that self-referential data cannot cause +// unbounded recursion. func deepCopy(src interface{}) interface{} { + return deepCopyValue(src, make(map[uintptr]interface{}), 0) +} + +const maxCopyDepth = 64 + +func deepCopyValue(src interface{}, seen map[uintptr]interface{}, depth int) interface{} { if src == nil { return nil } - // Handle basic types + // Fast path for immutable scalars. switch v := src.(type) { - case bool, int, int8, int16, int32, int64, uint, uint8, uint16, uint32, uint64, float32, float64, string: + case bool, int, int8, int16, int32, int64, + uint, uint8, uint16, uint32, uint64, + float32, float64, complex64, complex128, string: + return v + case time.Time: return v case []byte: dst := make([]byte, len(v)) @@ -295,54 +474,125 @@ func deepCopy(src interface{}) interface{} { return dst } - // Handle complex types using reflection + if depth >= maxCopyDepth { + return src + } + srcVal := reflect.ValueOf(src) - dstVal := reflect.New(srcVal.Type()).Elem() + switch srcVal.Kind() { + case reflect.Map: + if srcVal.IsNil() { + return src + } + if ptr := srcVal.Pointer(); ptr != 0 { + if existing, ok := seen[ptr]; ok { + return existing + } + } + dst := reflect.MakeMapWithSize(srcVal.Type(), srcVal.Len()) + if ptr := srcVal.Pointer(); ptr != 0 { + seen[ptr] = dst.Interface() + } + iter := srcVal.MapRange() + for iter.Next() { + copied := deepCopyValue(iter.Value().Interface(), seen, depth+1) + dst.SetMapIndex(iter.Key(), reflectValueFor(copied, iter.Value().Type())) + } + return dst.Interface() - deepCopyRecursive(srcVal, dstVal) - return dstVal.Interface() -} + case reflect.Slice: + if srcVal.IsNil() { + return src + } + dst := reflect.MakeSlice(srcVal.Type(), srcVal.Len(), srcVal.Len()) + for i := 0; i < srcVal.Len(); i++ { + copied := deepCopyValue(srcVal.Index(i).Interface(), seen, depth+1) + dst.Index(i).Set(reflectValueFor(copied, srcVal.Type().Elem())) + } + return dst.Interface() -// deepCopyRecursive performs recursive deep copying -func deepCopyRecursive(src, dst reflect.Value) { - switch src.Kind() { - case reflect.Pointer: - if src.IsNil() { - return + case reflect.Array: + dst := reflect.New(srcVal.Type()).Elem() + for i := 0; i < srcVal.Len(); i++ { + copied := deepCopyValue(srcVal.Index(i).Interface(), seen, depth+1) + dst.Index(i).Set(reflectValueFor(copied, srcVal.Type().Elem())) } - dst.Set(reflect.New(src.Type().Elem())) - deepCopyRecursive(src.Elem(), dst.Elem()) - case reflect.Interface: - if src.IsNil() { - return + return dst.Interface() + + case reflect.Ptr: + if srcVal.IsNil() { + return src } - dst.Set(reflect.ValueOf(deepCopy(src.Interface()))) - case reflect.Struct: - for i := 0; i < src.NumField(); i++ { - deepCopyRecursive(src.Field(i), dst.Field(i)) + elemType := srcVal.Type().Elem() + if !isCopyableStruct(elemType) { + // Cannot safely copy: share the pointer. + return src } - case reflect.Slice: - if src.IsNil() { - return + if ptr := srcVal.Pointer(); ptr != 0 { + if existing, ok := seen[ptr]; ok { + return existing + } } - dst.Set(reflect.MakeSlice(src.Type(), src.Len(), src.Cap())) - for i := 0; i < src.Len(); i++ { - deepCopyRecursive(src.Index(i), dst.Index(i)) + dst := reflect.New(elemType) + if ptr := srcVal.Pointer(); ptr != 0 { + seen[ptr] = dst.Interface() } - case reflect.Map: - if src.IsNil() { - return + copied := deepCopyValue(srcVal.Elem().Interface(), seen, depth+1) + dst.Elem().Set(reflectValueFor(copied, elemType)) + return dst.Interface() + + case reflect.Struct: + if !isCopyableStruct(srcVal.Type()) { + // Structs with unexported fields cannot be rebuilt via reflection. + // Returning the original preserves the value instead of panicking. + return src } - dst.Set(reflect.MakeMap(src.Type())) - for _, key := range src.MapKeys() { - srcVal := src.MapIndex(key) - dstVal := reflect.New(srcVal.Type()).Elem() - deepCopyRecursive(srcVal, dstVal) - dst.SetMapIndex(key, dstVal) + dst := reflect.New(srcVal.Type()).Elem() + for i := 0; i < srcVal.NumField(); i++ { + field := srcVal.Field(i) + if !dst.Field(i).CanSet() { + continue + } + copied := deepCopyValue(field.Interface(), seen, depth+1) + dst.Field(i).Set(reflectValueFor(copied, field.Type())) } + return dst.Interface() + default: - dst.Set(src) + // Chan, Func, UnsafePointer and scalars of named types: share as-is. + return src + } +} + +// reflectValueFor converts a copied value back into a reflect.Value assignable +// to the destination type, falling back to the zero value when the copy did not +// preserve assignability. +func reflectValueFor(v interface{}, dstType reflect.Type) reflect.Value { + if v == nil { + return reflect.Zero(dstType) + } + rv := reflect.ValueOf(v) + if rv.Type().AssignableTo(dstType) { + return rv + } + if rv.Type().ConvertibleTo(dstType) { + return rv.Convert(dstType) + } + return reflect.Zero(dstType) +} + +// isCopyableStruct reports whether every field of a struct type is exported, so +// that reflection can rebuild it field by field. +func isCopyableStruct(t reflect.Type) bool { + if t.Kind() != reflect.Struct { + return true + } + for i := 0; i < t.NumField(); i++ { + if t.Field(i).PkgPath != "" { // unexported + return false + } } + return true } // StateManager manages multiple states and provides advanced operations diff --git a/pkg/core/subgraph.go b/pkg/core/subgraph.go new file mode 100644 index 0000000..3e437d4 --- /dev/null +++ b/pkg/core/subgraph.go @@ -0,0 +1,167 @@ +// Copyright (c) 2024 GoLangGraph Team +// +// Licensed under the MIT License. See LICENSE file in the project root for full license information. +// +// Package: GoLangGraph - A powerful Go framework for building AI agent workflows + +package core + +import ( + "context" + "errors" + "fmt" +) + +// SubgraphOptions controls how a nested graph exchanges state with its parent. +type SubgraphOptions struct { + // InputKeys restricts what the subgraph sees. Empty means the whole state. + InputKeys []string + // OutputKeys restricts what is written back to the parent. Empty means the + // whole subgraph output state. + OutputKeys []string + // Namespace, when set, writes the subgraph output under this single key as + // a map[string]interface{} instead of merging keys into the parent state. + Namespace string + // Schema applies reducers when merging subgraph output into parent state. + // When nil, the parent graph's schema is used. + Schema *StateSchema + // PropagateInterrupts surfaces a subgraph interrupt to the parent instead of + // treating it as an error. + PropagateInterrupts bool +} + +// ErrSubgraphInterrupted wraps an interrupt raised inside a nested graph. +var ErrSubgraphInterrupted = errors.New("subgraph interrupted") + +// AddSubgraph registers a compiled graph as a node of this graph, mirroring +// LangGraph's ability to use a compiled graph as a node. +// +// The subgraph runs to completion with its own recursion limit and routing; its +// resulting state is merged back into the parent according to opts. +func (g *Graph) AddSubgraph(id, name string, sub *Graph, opts *SubgraphOptions) (*Node, error) { + if sub == nil { + return nil, fmt.Errorf("subgraph %s: nested graph must not be nil", id) + } + if sub == g { + return nil, fmt.Errorf("subgraph %s: a graph cannot contain itself", id) + } + if err := sub.Validate(); err != nil { + return nil, fmt.Errorf("subgraph %s: %w", id, err) + } + if contains, path := graphContains(sub, g, make(map[*Graph]bool)); contains { + return nil, fmt.Errorf("subgraph %s would create a cycle of graphs: %s", id, path) + } + + if opts == nil { + opts = &SubgraphOptions{} + } + // Copy so later caller mutations cannot change node behavior. + local := *opts + local.InputKeys = append([]string(nil), opts.InputKeys...) + local.OutputKeys = append([]string(nil), opts.OutputKeys...) + + fn := func(ctx context.Context, state *BaseState) (*BaseState, error) { + input := projectState(state, local.InputKeys) + + out, err := sub.Execute(ctx, input) + if err != nil { + var ie *InterruptError + if errors.As(err, &ie) && local.PropagateInterrupts { + return nil, fmt.Errorf("%w: %s: %w", ErrSubgraphInterrupted, id, err) + } + return nil, fmt.Errorf("subgraph %s failed: %w", id, err) + } + + merged := state.Clone() + schema := local.Schema + if schema == nil { + schema = g.StateSchema() + } + + if local.Namespace != "" { + merged.Set(local.Namespace, out.GetAll()) + return merged, nil + } + + projected := projectState(out, local.OutputKeys) + merged.MergeWithSchema(projected, schema) + return merged, nil + } + + node := g.AddNode(id, name, fn) + node.Metadata["type"] = "subgraph" + node.Metadata["subgraph_id"] = sub.ID + node.Metadata["subgraph_name"] = sub.Name + + g.mu.Lock() + if g.subgraphs == nil { + g.subgraphs = make(map[string]*Graph) + } + g.subgraphs[id] = sub + g.mu.Unlock() + + return node, nil +} + +// Subgraph returns the nested graph registered under a node ID. +func (g *Graph) Subgraph(nodeID string) (*Graph, bool) { + g.mu.RLock() + defer g.mu.RUnlock() + sub, ok := g.subgraphs[nodeID] + return sub, ok +} + +// Subgraphs returns nested graphs by node ID. +func (g *Graph) Subgraphs() map[string]*Graph { + g.mu.RLock() + defer g.mu.RUnlock() + out := make(map[string]*Graph, len(g.subgraphs)) + for k, v := range g.subgraphs { + out[k] = v + } + return out +} + +// graphContains reports whether target is reachable from root through nested +// subgraphs, which would make the composition infinitely recursive. +func graphContains(root, target *Graph, seen map[*Graph]bool) (bool, string) { + if root == nil || seen[root] { + return false, "" + } + seen[root] = true + + root.mu.RLock() + nested := make(map[string]*Graph, len(root.subgraphs)) + for k, v := range root.subgraphs { + nested[k] = v + } + root.mu.RUnlock() + + for id, sub := range nested { + if sub == target { + return true, fmt.Sprintf("%s -> %s", root.Name, id) + } + if found, path := graphContains(sub, target, seen); found { + return true, fmt.Sprintf("%s -> %s", root.Name, path) + } + } + return false, "" +} + +// projectState returns a state limited to the given keys. Empty keys means the +// whole state. +func projectState(state *BaseState, keys []string) *BaseState { + if state == nil { + return NewBaseState() + } + if len(keys) == 0 { + return state.Clone() + } + out := NewBaseState() + for _, k := range keys { + if v, ok := state.Get(k); ok { + out.Set(k, v) + } + } + return out +} diff --git a/pkg/debug/visualizer.go b/pkg/debug/visualizer.go index edd1831..731024d 100644 --- a/pkg/debug/visualizer.go +++ b/pkg/debug/visualizer.go @@ -377,8 +377,14 @@ func (gv *GraphVisualizer) getConditionName(edge *core.Edge) string { if edge.Condition == nil { return "" } - // This is a placeholder - in a real implementation, you'd want to - // extract meaningful condition names from the condition function + // A Go function value carries no name at runtime, so the label is generic. + // Attach a "condition_name" entry to the edge metadata for something more + // descriptive in the rendered graph. + if edge.Metadata != nil { + if name, ok := edge.Metadata["condition_name"].(string); ok && name != "" { + return name + } + } return "condition" } diff --git a/pkg/llm/early_exit.go b/pkg/llm/early_exit.go index 85e001b..0c82304 100644 --- a/pkg/llm/early_exit.go +++ b/pkg/llm/early_exit.go @@ -19,7 +19,7 @@ import ( // ErrStreamEarlyExit is returned by a StreamCallback (or CollectStream) when a // complete tool-call / structured JSON result is already formed and the rest of -// the token stream should be cancelled to save decode latency on SLMs. +// the token stream should be canceled to save decode latency on SLMs. var ErrStreamEarlyExit = errors.New("llm: stream early exit") // IsStreamEarlyExit reports whether err is (or wraps) ErrStreamEarlyExit. @@ -184,7 +184,7 @@ func CollectStream( } resp.Metadata["early_exit"] = true } - // Providers often omit usage on cancelled streams β€” estimate so callers can account. + // Providers often omit usage on canceled streams β€” estimate so callers can account. _ = EnsureUsage(resp, req) if resp == nil || (resp.Choices[0].Message.Content == "" && len(resp.Choices[0].Message.ToolCalls) == 0) { if err != nil && !IsStreamEarlyExit(err) { diff --git a/pkg/llm/gemini.go b/pkg/llm/gemini.go index bc52d9a..19ad00d 100644 --- a/pkg/llm/gemini.go +++ b/pkg/llm/gemini.go @@ -7,38 +7,244 @@ package llm import ( + "bufio" + "bytes" "context" + "encoding/json" "fmt" + "io" + "net/http" + "net/url" "strings" "time" "github.com/sirupsen/logrus" ) -// GeminiProvider implements the Provider interface for Google Gemini -// This is a mock implementation for demonstration purposes +// DefaultGeminiEndpoint is the Generative Language API base URL. +const DefaultGeminiEndpoint = "https://generativelanguage.googleapis.com/v1beta" + +// GeminiProvider implements the Provider interface against Google's Generative +// Language API. +// +// This provider previously returned hardcoded strings β€” "Hello! I'm Gemini..." +// and a note that a real implementation would call the API β€” while presenting +// itself as a working provider. Configuring it with a valid API key produced +// canned text with nothing to indicate the model had never been contacted. type GeminiProvider struct { config *ProviderConfig logger *logrus.Logger models []string - lastSync time.Time + client *http.Client + endpoint string } // NewGeminiProvider creates a new Gemini provider func NewGeminiProvider(config *ProviderConfig) (*GeminiProvider, error) { + if config == nil { + return nil, fmt.Errorf("gemini configuration is required") + } if config.APIKey == "" { - return nil, fmt.Errorf("Gemini API key is required") + return nil, fmt.Errorf("gemini API key is required") + } + + timeout := config.Timeout + if timeout <= 0 { + timeout = 60 * time.Second + } + + endpoint := strings.TrimRight(config.Endpoint, "/") + if endpoint == "" { + endpoint = DefaultGeminiEndpoint } provider := &GeminiProvider{ - config: config, - logger: logrus.New(), - models: []string{"gemini-pro", "gemini-pro-vision"}, + config: config, + logger: logrus.New(), + models: []string{"gemini-1.5-pro", "gemini-1.5-flash", "gemini-pro"}, + client: &http.Client{Timeout: timeout}, + endpoint: endpoint, } return provider, nil } +// --------------------------------------------------------------------------- +// Generative Language API wire types +// --------------------------------------------------------------------------- + +type geminiPart struct { + Text string `json:"text,omitempty"` +} + +type geminiContent struct { + Role string `json:"role,omitempty"` + Parts []geminiPart `json:"parts"` +} + +type geminiGenerationConfig struct { + Temperature *float64 `json:"temperature,omitempty"` + MaxOutputTokens *int `json:"maxOutputTokens,omitempty"` + TopP *float64 `json:"topP,omitempty"` + StopSequences []string `json:"stopSequences,omitempty"` +} + +type geminiRequest struct { + Contents []geminiContent `json:"contents"` + SystemInstruction *geminiContent `json:"systemInstruction,omitempty"` + GenerationConfig *geminiGenerationConfig `json:"generationConfig,omitempty"` +} + +type geminiCandidate struct { + Content geminiContent `json:"content"` + FinishReason string `json:"finishReason"` + Index int `json:"index"` +} + +type geminiUsage struct { + PromptTokenCount int `json:"promptTokenCount"` + CandidatesTokenCount int `json:"candidatesTokenCount"` + TotalTokenCount int `json:"totalTokenCount"` +} + +type geminiResponse struct { + Candidates []geminiCandidate `json:"candidates"` + UsageMetadata geminiUsage `json:"usageMetadata"` + Error *struct { + Code int `json:"code"` + Message string `json:"message"` + Status string `json:"status"` + } `json:"error,omitempty"` +} + +// modelFor resolves the model for a request, falling back to configuration. +func (p *GeminiProvider) modelFor(req CompletionRequest) string { + if req.Model != "" { + return req.Model + } + if p.config.Model != "" { + return p.config.Model + } + return "gemini-1.5-flash" +} + +// buildRequest converts a CompletionRequest into the Gemini wire format. +// +// Gemini names the assistant role "model" and carries system prompts in a +// dedicated systemInstruction field rather than as a message. +func (p *GeminiProvider) buildRequest(req CompletionRequest) (*geminiRequest, error) { + if len(req.Messages) == 0 { + return nil, fmt.Errorf("no messages provided") + } + + out := &geminiRequest{} + var systemParts []geminiPart + + if req.SystemPrompt != "" { + systemParts = append(systemParts, geminiPart{Text: req.SystemPrompt}) + } + + for _, msg := range req.Messages { + switch msg.Role { + case "system": + systemParts = append(systemParts, geminiPart{Text: msg.Content}) + case "assistant": + out.Contents = append(out.Contents, geminiContent{ + Role: "model", Parts: []geminiPart{{Text: msg.Content}}, + }) + default: + out.Contents = append(out.Contents, geminiContent{ + Role: "user", Parts: []geminiPart{{Text: msg.Content}}, + }) + } + } + + if len(out.Contents) == 0 { + return nil, fmt.Errorf("no user or assistant messages provided") + } + if len(systemParts) > 0 { + out.SystemInstruction = &geminiContent{Parts: systemParts} + } + + cfg := &geminiGenerationConfig{} + set := false + if req.Temperature != 0 { + t := req.Temperature + cfg.Temperature = &t + set = true + } else if p.config.Temperature != 0 { + t := p.config.Temperature + cfg.Temperature = &t + set = true + } + if req.MaxTokens > 0 { + m := req.MaxTokens + cfg.MaxOutputTokens = &m + set = true + } else if p.config.MaxTokens > 0 { + m := p.config.MaxTokens + cfg.MaxOutputTokens = &m + set = true + } + if len(req.StopSequences) > 0 { + cfg.StopSequences = req.StopSequences + set = true + } + if set { + out.GenerationConfig = cfg + } + + return out, nil +} + +// callURL builds an API URL with the key supplied as a query parameter, which +// is how the Generative Language API authenticates. +func (p *GeminiProvider) callURL(model, method string, extra url.Values) string { + values := url.Values{} + for k, v := range extra { + values[k] = v + } + values.Set("key", p.config.APIKey) + return fmt.Sprintf("%s/models/%s:%s?%s", p.endpoint, url.PathEscape(model), method, values.Encode()) +} + +// toCompletionResponse converts a Gemini response into the common format. +func toCompletionResponse(model string, resp *geminiResponse) *CompletionResponse { + out := &CompletionResponse{ + ID: fmt.Sprintf("gemini-%d", time.Now().UnixNano()), + Object: "chat.completion", + Created: time.Now().Unix(), + Model: model, + Usage: Usage{ + PromptTokens: resp.UsageMetadata.PromptTokenCount, + CompletionTokens: resp.UsageMetadata.CandidatesTokenCount, + TotalTokens: resp.UsageMetadata.TotalTokenCount, + }, + } + + for _, candidate := range resp.Candidates { + var text strings.Builder + for _, part := range candidate.Content.Parts { + text.WriteString(part.Text) + } + out.Choices = append(out.Choices, Choice{ + Index: candidate.Index, + Message: Message{Role: "assistant", Content: text.String()}, + FinishReason: strings.ToLower(candidate.FinishReason), + }) + } + + if len(out.Choices) == 0 { + // A response with no candidates usually means the prompt was blocked. + out.Choices = append(out.Choices, Choice{ + Index: 0, + Message: Message{Role: "assistant", Content: ""}, + FinishReason: "stop", + }) + } + return out +} + // GetName returns the provider name func (p *GeminiProvider) GetName() string { return "gemini" @@ -49,94 +255,183 @@ func (p *GeminiProvider) GetModels(ctx context.Context) ([]string, error) { return p.models, nil } -// Complete generates a completion +// Complete generates a completion by calling the Generative Language API. func (p *GeminiProvider) Complete(ctx context.Context, req CompletionRequest) (*CompletionResponse, error) { - // Mock implementation - in a real implementation, this would call the Gemini API - if len(req.Messages) == 0 { - return nil, fmt.Errorf("no messages provided") + payload, err := p.buildRequest(req) + if err != nil { + return nil, err } + body, err := json.Marshal(payload) + if err != nil { + return nil, fmt.Errorf("failed to marshal request: %w", err) + } + + model := p.modelFor(req) + var result *CompletionResponse + + attempt := func(ctx context.Context) error { + httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, + p.callURL(model, "generateContent", nil), bytes.NewReader(body)) + if err != nil { + return fmt.Errorf("failed to create request: %w", err) + } + httpReq.Header.Set("Content-Type", "application/json") + + resp, err := p.client.Do(httpReq) + if err != nil { + if ctxErr := ctx.Err(); ctxErr != nil { + return ctxErr + } + return NewTransportError("Gemini", err) + } + defer func() { _ = resp.Body.Close() }() - lastMessage := req.Messages[len(req.Messages)-1] + if resp.StatusCode != http.StatusOK { + return NewProviderError("Gemini", resp) + } - // Generate a mock response based on the input - var responseText string - if strings.Contains(strings.ToLower(lastMessage.Content), "hello") { - responseText = "Hello! I'm Gemini, Google's AI assistant. How can I help you today?" - } else if strings.Contains(strings.ToLower(lastMessage.Content), "go programming") { - responseText = "Go is a fantastic programming language! It's known for its simplicity, excellent concurrency support with goroutines, and strong performance. It's perfect for building scalable backend services, CLI tools, and distributed systems." - } else { - responseText = "I understand your request. This is a mock Gemini response for demonstration purposes. In a real implementation, this would be powered by Google's Gemini API." + raw, err := io.ReadAll(limitedBody(resp.Body)) + if err != nil { + return fmt.Errorf("failed to read response: %w", err) + } + + var decoded geminiResponse + if err := json.Unmarshal(raw, &decoded); err != nil { + return fmt.Errorf("failed to decode response: %w", err) + } + if decoded.Error != nil { + return fmt.Errorf("%w: gemini API error: %s", ErrProviderRequest, decoded.Error.Message) + } + + result = toCompletionResponse(model, &decoded) + return nil } - return &CompletionResponse{ - ID: fmt.Sprintf("gemini-mock-%d", time.Now().Unix()), - Object: "chat.completion", - Created: time.Now().Unix(), - Model: req.Model, - Choices: []Choice{ - { - Index: 0, - Message: Message{ - Role: "assistant", - Content: responseText, - }, - FinishReason: "stop", - }, - }, - Usage: Usage{ - PromptTokens: len(lastMessage.Content) / 4, - CompletionTokens: len(responseText) / 4, - TotalTokens: (len(lastMessage.Content) + len(responseText)) / 4, - }, - }, nil + if err := WithRetry(ctx, p.config, attempt); err != nil { + return nil, err + } + return result, nil } -// CompleteStream generates a streaming completion +// CompleteStream streams a completion using the API's server-sent events +// endpoint, invoking the callback for each chunk as it arrives. func (p *GeminiProvider) CompleteStream(ctx context.Context, req CompletionRequest, callback StreamCallback) error { - // Mock streaming implementation - response, err := p.Complete(ctx, req) + payload, err := p.buildRequest(req) if err != nil { return err } + body, err := json.Marshal(payload) + if err != nil { + return fmt.Errorf("failed to marshal request: %w", err) + } - // Simulate streaming by sending the response in chunks - content := response.Choices[0].Message.Content - words := strings.Fields(content) - - for i, word := range words { - chunk := CompletionResponse{ - ID: fmt.Sprintf("gemini-stream-%d", i), - Object: "chat.completion.chunk", - Created: time.Now().Unix(), - Model: req.Model, - Choices: []Choice{ - { - Index: 0, - Delta: Message{ - Role: "assistant", - Content: word + " ", - }, - }, - }, + model := p.modelFor(req) + target := p.callURL(model, "streamGenerateContent", url.Values{"alt": []string{"sse"}}) + + httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, target, bytes.NewReader(body)) + if err != nil { + return fmt.Errorf("failed to create request: %w", err) + } + httpReq.Header.Set("Content-Type", "application/json") + httpReq.Header.Set("Accept", "text/event-stream") + + resp, err := p.client.Do(httpReq) + if err != nil { + if ctxErr := ctx.Err(); ctxErr != nil { + return ctxErr } + return NewTransportError("Gemini", err) + } + defer func() { _ = resp.Body.Close() }() - if err := callback(chunk); err != nil { - if IsStreamEarlyExit(err) { - return err - } + if resp.StatusCode != http.StatusOK { + return NewProviderError("Gemini", resp) + } + + scanner := bufio.NewScanner(io.LimitReader(resp.Body, MaxResponseBytes)) + scanner.Buffer(make([]byte, 0, 64*1024), 4*1024*1024) + + index := 0 + for scanner.Scan() { + if err := ctx.Err(); err != nil { return err } - // Small delay to simulate streaming - time.Sleep(50 * time.Millisecond) + line := strings.TrimSpace(scanner.Text()) + if line == "" || !strings.HasPrefix(line, "data:") { + continue + } + data := strings.TrimSpace(strings.TrimPrefix(line, "data:")) + if data == "" || data == "[DONE]" { + continue + } + + var decoded geminiResponse + if err := json.Unmarshal([]byte(data), &decoded); err != nil { + return fmt.Errorf("failed to decode stream chunk: %w", err) + } + if decoded.Error != nil { + return fmt.Errorf("%w: gemini API error: %s", ErrProviderRequest, decoded.Error.Message) + } + + for _, candidate := range decoded.Candidates { + var text strings.Builder + for _, part := range candidate.Content.Parts { + text.WriteString(part.Text) + } + chunk := CompletionResponse{ + ID: fmt.Sprintf("gemini-stream-%d", index), + Object: "chat.completion.chunk", + Created: time.Now().Unix(), + Model: model, + Choices: []Choice{{ + Index: candidate.Index, + Delta: Message{Role: "assistant", Content: text.String()}, + FinishReason: strings.ToLower(candidate.FinishReason), + }}, + Usage: Usage{ + PromptTokens: decoded.UsageMetadata.PromptTokenCount, + CompletionTokens: decoded.UsageMetadata.CandidatesTokenCount, + TotalTokens: decoded.UsageMetadata.TotalTokenCount, + }, + } + if err := callback(chunk); err != nil { + if IsStreamEarlyExit(err) { + return err + } + return fmt.Errorf("callback error: %w", err) + } + index++ + } } + if err := scanner.Err(); err != nil { + if ctxErr := ctx.Err(); ctxErr != nil { + return ctxErr + } + return fmt.Errorf("failed to read stream: %w", err) + } return nil } // IsHealthy checks if the provider is healthy func (p *GeminiProvider) IsHealthy(ctx context.Context) error { - // Mock health check - always healthy for demonstration + target := fmt.Sprintf("%s/models?key=%s", p.endpoint, url.QueryEscape(p.config.APIKey)) + + req, err := http.NewRequestWithContext(ctx, http.MethodGet, target, nil) + if err != nil { + return fmt.Errorf("failed to create health request: %w", err) + } + + resp, err := p.client.Do(req) + if err != nil { + return NewTransportError("Gemini", err) + } + defer func() { _ = resp.Body.Close() }() + + if resp.StatusCode != http.StatusOK { + return NewProviderError("Gemini", resp) + } return nil } diff --git a/pkg/llm/gemini_test.go b/pkg/llm/gemini_test.go new file mode 100644 index 0000000..2f2eecf --- /dev/null +++ b/pkg/llm/gemini_test.go @@ -0,0 +1,404 @@ +// Copyright (c) 2024 GoLangGraph Team +// +// Licensed under the MIT License. See LICENSE file in the project root for full license information. + +package llm + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" + "net/http/httptest" + "strings" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// geminiAPI is a stand-in for Google's Generative Language API. The provider +// under test is the real one: it builds real requests, speaks the real wire +// format and parses real responses. +type geminiAPI struct { + server *httptest.Server + requests atomic.Int32 + + // lastBody is the decoded body of the most recent generateContent call. + lastBody map[string]interface{} + lastPath string + lastKey string +} + +func newGeminiAPI(t *testing.T, handler http.HandlerFunc) *geminiAPI { + t.Helper() + api := &geminiAPI{} + api.server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + api.requests.Add(1) + api.lastPath = r.URL.Path + api.lastKey = r.URL.Query().Get("key") + + if r.Body != nil { + if raw, err := io.ReadAll(r.Body); err == nil && len(raw) > 0 { + var decoded map[string]interface{} + if json.Unmarshal(raw, &decoded) == nil { + api.lastBody = decoded + } + } + } + handler(w, r) + })) + t.Cleanup(api.server.Close) + return api +} + +// provider builds a Gemini provider pointed at the fake API. +func (api *geminiAPI) provider(t *testing.T, mutate func(*ProviderConfig)) *GeminiProvider { + t.Helper() + cfg := DefaultProviderConfig() + cfg.Type = "gemini" + cfg.APIKey = "test-key" // pragma: allowlist secret + cfg.Model = "gemini-1.5-flash" + cfg.Endpoint = api.server.URL + cfg.RetryCount = 0 + cfg.RetryDelay = time.Millisecond + cfg.Timeout = 5 * time.Second + if mutate != nil { + mutate(cfg) + } + p, err := NewGeminiProvider(cfg) + require.NoError(t, err) + return p +} + +func geminiSuccess(text string) http.HandlerFunc { + return func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(map[string]interface{}{ + "candidates": []map[string]interface{}{{ + "content": map[string]interface{}{"role": "model", "parts": []map[string]string{{"text": text}}}, + "finishReason": "STOP", + "index": 0, + }}, + "usageMetadata": map[string]int{ + "promptTokenCount": 7, "candidatesTokenCount": 11, "totalTokenCount": 18, + }, + }) + } +} + +func geminiTestRequest() CompletionRequest { + return CompletionRequest{ + Messages: []Message{{Role: "user", Content: "Hello, how are you?"}}, + } +} + +// The provider must call the model rather than returning canned text. It +// previously returned hardcoded strings without contacting the API at all. +func TestGemini_CallsTheAPI(t *testing.T) { + api := newGeminiAPI(t, geminiSuccess("a genuine model reply")) + p := api.provider(t, nil) + + resp, err := p.Complete(context.Background(), geminiTestRequest()) + require.NoError(t, err) + + assert.EqualValues(t, 1, api.requests.Load(), "the provider must actually call the API") + require.NotEmpty(t, resp.Choices) + assert.Equal(t, "a genuine model reply", resp.Choices[0].Message.Content, + "the reply must come from the API, not from the provider") + assert.Equal(t, "stop", resp.Choices[0].FinishReason) + assert.Equal(t, 7, resp.Usage.PromptTokens) + assert.Equal(t, 11, resp.Usage.CompletionTokens) + assert.Equal(t, 18, resp.Usage.TotalTokens) +} + +// The request must target the configured model and carry the API key. +func TestGemini_RequestTargetsModelAndKey(t *testing.T) { + api := newGeminiAPI(t, geminiSuccess("ok")) + p := api.provider(t, func(c *ProviderConfig) { c.Model = "gemini-1.5-pro" }) + + _, err := p.Complete(context.Background(), geminiTestRequest()) + require.NoError(t, err) + + assert.Contains(t, api.lastPath, "gemini-1.5-pro") + assert.Contains(t, api.lastPath, "generateContent") + assert.Equal(t, "test-key", api.lastKey) +} + +// Gemini names the assistant role "model" and carries the system prompt in a +// dedicated field rather than as a message. +func TestGemini_MessageMapping(t *testing.T) { + api := newGeminiAPI(t, geminiSuccess("ok")) + p := api.provider(t, nil) + + _, err := p.Complete(context.Background(), CompletionRequest{ + SystemPrompt: "You are terse.", + Messages: []Message{ + {Role: "system", Content: "Also be polite."}, + {Role: "user", Content: "hi"}, + {Role: "assistant", Content: "hello"}, + {Role: "user", Content: "again"}, + }, + }) + require.NoError(t, err) + + body := api.lastBody + require.NotNil(t, body) + + system, ok := body["systemInstruction"].(map[string]interface{}) + require.True(t, ok, "system prompts belong in systemInstruction: %v", body) + assert.Contains(t, fmt.Sprint(system), "You are terse.") + assert.Contains(t, fmt.Sprint(system), "Also be polite.") + + contents, ok := body["contents"].([]interface{}) + require.True(t, ok) + require.Len(t, contents, 3, "only user and assistant turns become contents") + + roles := make([]string, 0, len(contents)) + for _, c := range contents { + roles = append(roles, fmt.Sprint(c.(map[string]interface{})["role"])) + } + assert.Equal(t, []string{"user", "model", "user"}, roles, + "the assistant role must be sent as \"model\"") +} + +// Generation settings must reach the API. +func TestGemini_GenerationConfig(t *testing.T) { + api := newGeminiAPI(t, geminiSuccess("ok")) + p := api.provider(t, nil) + + _, err := p.Complete(context.Background(), CompletionRequest{ + Messages: []Message{{Role: "user", Content: "hi"}}, + Temperature: 0.25, + MaxTokens: 512, + StopSequences: []string{"END"}, + }) + require.NoError(t, err) + + cfg, ok := api.lastBody["generationConfig"].(map[string]interface{}) + require.True(t, ok, "generation settings must be sent: %v", api.lastBody) + assert.InDelta(t, 0.25, cfg["temperature"], 1e-9) + assert.EqualValues(t, 512, cfg["maxOutputTokens"]) + assert.Contains(t, fmt.Sprint(cfg["stopSequences"]), "END") +} + +func TestGemini_EmptyMessagesRejected(t *testing.T) { + api := newGeminiAPI(t, geminiSuccess("ok")) + p := api.provider(t, nil) + + _, err := p.Complete(context.Background(), CompletionRequest{}) + require.Error(t, err) + assert.Zero(t, api.requests.Load(), "an invalid request must not reach the API") +} + +// API errors must be classified like any other provider's. +func TestGemini_ErrorClassification(t *testing.T) { + cases := []struct { + status int + want error + }{ + {http.StatusTooManyRequests, ErrRateLimited}, + {http.StatusUnauthorized, ErrProviderAuth}, + {http.StatusBadRequest, ErrProviderRequest}, + {http.StatusInternalServerError, ErrProviderUnavailable}, + } + + for _, tc := range cases { + t.Run(fmt.Sprint(tc.status), func(t *testing.T) { + api := newGeminiAPI(t, func(w http.ResponseWriter, r *http.Request) { + http.Error(w, "upstream said no", tc.status) + }) + p := api.provider(t, nil) + + _, err := p.Complete(context.Background(), geminiTestRequest()) + require.Error(t, err) + assert.True(t, errors.Is(err, tc.want), "status %d: got %v", tc.status, err) + }) + } +} + +// An error carried inside a 200 body is a permanent request error. +func TestGemini_InBodyErrorIsReported(t *testing.T) { + api := newGeminiAPI(t, func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(map[string]interface{}{ + "error": map[string]interface{}{"code": 400, "message": "API key not valid", "status": "INVALID_ARGUMENT"}, + }) + }) + p := api.provider(t, nil) + + _, err := p.Complete(context.Background(), geminiTestRequest()) + require.Error(t, err) + assert.Contains(t, err.Error(), "API key not valid") + assert.True(t, errors.Is(err, ErrProviderRequest)) +} + +func TestGemini_MalformedResponse(t *testing.T) { + api := newGeminiAPI(t, func(w http.ResponseWriter, r *http.Request) { + _, _ = w.Write([]byte("{not json")) + }) + p := api.provider(t, nil) + + _, err := p.Complete(context.Background(), geminiTestRequest()) + require.Error(t, err) + assert.Contains(t, err.Error(), "decode") +} + +// A blocked prompt returns no candidates; that must be a usable response +// rather than an index-out-of-range panic. +func TestGemini_NoCandidates(t *testing.T) { + api := newGeminiAPI(t, func(w http.ResponseWriter, r *http.Request) { + _ = json.NewEncoder(w).Encode(map[string]interface{}{"candidates": []interface{}{}}) + }) + p := api.provider(t, nil) + + resp, err := p.Complete(context.Background(), geminiTestRequest()) + require.NoError(t, err) + require.NotEmpty(t, resp.Choices) + assert.Empty(t, resp.Choices[0].Message.Content) +} + +func TestGemini_RetriesTransientFailures(t *testing.T) { + var calls atomic.Int32 + api := newGeminiAPI(t, func(w http.ResponseWriter, r *http.Request) { + if calls.Add(1) < 3 { + http.Error(w, "unavailable", http.StatusServiceUnavailable) + return + } + geminiSuccess("recovered")(w, r) + }) + p := api.provider(t, func(c *ProviderConfig) { c.RetryCount = 3 }) + + resp, err := p.Complete(context.Background(), geminiTestRequest()) + require.NoError(t, err) + assert.Equal(t, "recovered", resp.Choices[0].Message.Content) +} + +func TestGemini_Cancellation(t *testing.T) { + release := make(chan struct{}) + api := newGeminiAPI(t, func(w http.ResponseWriter, r *http.Request) { + <-release + }) + defer close(release) + + p := api.provider(t, func(c *ProviderConfig) { c.Timeout = 30 * time.Second }) + + ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond) + defer cancel() + + start := time.Now() + _, err := p.Complete(ctx, geminiTestRequest()) + require.Error(t, err) + assert.Less(t, time.Since(start), 10*time.Second) +} + +// Streaming must consume the server-sent event stream and deliver each chunk. +func TestGemini_Streaming(t *testing.T) { + api := newGeminiAPI(t, func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + flusher, _ := w.(http.Flusher) + for _, text := range []string{"Hello", " there", "!"} { + frame, _ := json.Marshal(map[string]interface{}{ + "candidates": []map[string]interface{}{{ + "content": map[string]interface{}{"role": "model", "parts": []map[string]string{{"text": text}}}, + "index": 0, + }}, + }) + _, _ = fmt.Fprintf(w, "data: %s\n\n", frame) + if flusher != nil { + flusher.Flush() + } + } + }) + p := api.provider(t, nil) + + var chunks []string + err := p.CompleteStream(context.Background(), geminiTestRequest(), func(chunk CompletionResponse) error { + require.NotEmpty(t, chunk.Choices) + chunks = append(chunks, chunk.Choices[0].Delta.Content) + return nil + }) + require.NoError(t, err) + + assert.Equal(t, []string{"Hello", " there", "!"}, chunks) + assert.Contains(t, api.lastPath, "streamGenerateContent") + assert.Equal(t, strings.Join(chunks, ""), "Hello there!") +} + +// A callback that fails must stop the stream and surface its error. +func TestGemini_StreamingCallbackError(t *testing.T) { + api := newGeminiAPI(t, func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + for i := 0; i < 5; i++ { + frame, _ := json.Marshal(map[string]interface{}{ + "candidates": []map[string]interface{}{{ + "content": map[string]interface{}{"parts": []map[string]string{{"text": "x"}}}, + }}, + }) + _, _ = fmt.Fprintf(w, "data: %s\n\n", frame) + } + }) + p := api.provider(t, nil) + + sentinel := errors.New("consumer gave up") + err := p.CompleteStream(context.Background(), geminiTestRequest(), func(chunk CompletionResponse) error { + return sentinel + }) + assert.ErrorIs(t, err, sentinel) +} + +func TestGemini_StreamingErrorStatus(t *testing.T) { + api := newGeminiAPI(t, func(w http.ResponseWriter, r *http.Request) { + http.Error(w, "nope", http.StatusUnauthorized) + }) + p := api.provider(t, nil) + + err := p.CompleteStream(context.Background(), geminiTestRequest(), func(chunk CompletionResponse) error { + return nil + }) + require.Error(t, err) + assert.True(t, errors.Is(err, ErrProviderAuth)) +} + +// Health must reflect the API, not a hardcoded success. +func TestGemini_Health(t *testing.T) { + t.Run("healthy", func(t *testing.T) { + api := newGeminiAPI(t, func(w http.ResponseWriter, r *http.Request) { + _ = json.NewEncoder(w).Encode(map[string]interface{}{"models": []interface{}{}}) + }) + assert.NoError(t, api.provider(t, nil).IsHealthy(context.Background())) + }) + + t.Run("rejected credentials", func(t *testing.T) { + api := newGeminiAPI(t, func(w http.ResponseWriter, r *http.Request) { + http.Error(w, "bad key", http.StatusUnauthorized) + }) + err := api.provider(t, nil).IsHealthy(context.Background()) + require.Error(t, err, "an unhealthy provider must not report healthy") + assert.True(t, errors.Is(err, ErrProviderAuth)) + }) +} + +func TestGemini_RequiresAPIKey(t *testing.T) { + cfg := DefaultProviderConfig() + cfg.Type = "gemini" + _, err := NewGeminiProvider(cfg) + assert.Error(t, err) + + _, err = NewGeminiProvider(nil) + assert.Error(t, err) +} + +// Credentials must never appear in the configuration the API exposes. +func TestGemini_ConfigMasksAPIKey(t *testing.T) { + api := newGeminiAPI(t, geminiSuccess("ok")) + p := api.provider(t, nil) + + cfg := p.GetConfig() + assert.NotEqual(t, "test-key", cfg["api_key"]) + assert.NotContains(t, fmt.Sprint(cfg), "test-key") +} diff --git a/pkg/llm/ollama.go b/pkg/llm/ollama.go index b1930f4..fb2a3af 100644 --- a/pkg/llm/ollama.go +++ b/pkg/llm/ollama.go @@ -10,6 +10,7 @@ import ( "bytes" "context" "encoding/json" + "errors" "fmt" "io" "net/http" @@ -120,7 +121,7 @@ func (p *OllamaProvider) GetModels(ctx context.Context) ([]string, error) { if err != nil { return nil, fmt.Errorf("failed to get models: %w", err) } - defer resp.Body.Close() + defer func() { _ = resp.Body.Close() }() if resp.StatusCode != http.StatusOK { return nil, fmt.Errorf("failed to get models: status %d", resp.StatusCode) @@ -155,52 +156,68 @@ func (p *OllamaProvider) Complete(ctx context.Context, req CompletionRequest) (* "model": ollamaReq.Model, }).Debug("Sending request to Ollama") - httpReq, err := http.NewRequestWithContext(ctx, "POST", p.config.Endpoint+"/api/chat", bytes.NewBuffer(reqBody)) - if err != nil { - return nil, fmt.Errorf("failed to create request: %w", err) - } - - httpReq.Header.Set("Content-Type", "application/json") - - resp, err := p.client.Do(httpReq) - if err != nil { - return nil, fmt.Errorf("failed to send request: %w", err) - } - defer resp.Body.Close() - - if resp.StatusCode != http.StatusOK { - body, _ := io.ReadAll(resp.Body) - return nil, fmt.Errorf("Ollama API error: status %d, body: %s", resp.StatusCode, string(body)) - } - - // Read all streaming chunks until done=true var completeResponse strings.Builder var finalModel string var finalRole string - decoder := json.NewDecoder(resp.Body) - for { - var ollamaResp OllamaResponse - if err := decoder.Decode(&ollamaResp); err != nil { - if err == io.EOF { - break + // Transient failures (network errors, 5xx, rate limits) are retried + // according to the provider configuration; permanent errors fail fast. + attempt := func(ctx context.Context) error { + completeResponse.Reset() + finalModel = "" + finalRole = "" + + httpReq, err := http.NewRequestWithContext(ctx, "POST", p.config.Endpoint+"/api/chat", bytes.NewBuffer(reqBody)) + if err != nil { + return fmt.Errorf("failed to create request: %w", err) + } + httpReq.Header.Set("Content-Type", "application/json") + + resp, err := p.client.Do(httpReq) + if err != nil { + if ctxErr := ctx.Err(); ctxErr != nil { + return ctxErr } - return nil, fmt.Errorf("failed to decode response: %w", err) + return NewTransportError("Ollama", err) } + defer func() { _ = resp.Body.Close() }() - if ollamaResp.Error != "" { - return nil, fmt.Errorf("Ollama API error: %s", ollamaResp.Error) + if resp.StatusCode != http.StatusOK { + return NewProviderError("Ollama", resp) } - // Accumulate the response content - completeResponse.WriteString(ollamaResp.Message.Content) - finalModel = ollamaResp.Model - finalRole = ollamaResp.Message.Role + // Read all streaming chunks until done=true, bounded so a provider that + // never terminates cannot exhaust memory. + decoder := json.NewDecoder(limitedBody(resp.Body)) + for { + var ollamaResp OllamaResponse + if err := decoder.Decode(&ollamaResp); err != nil { + if errors.Is(err, io.EOF) { + break + } + if ctxErr := ctx.Err(); ctxErr != nil { + return ctxErr + } + return fmt.Errorf("failed to decode response: %w", err) + } + + if ollamaResp.Error != "" { + return fmt.Errorf("%w: Ollama API error: %s", ErrProviderRequest, ollamaResp.Error) + } - // Break when done - if ollamaResp.Done { - break + completeResponse.WriteString(ollamaResp.Message.Content) + finalModel = ollamaResp.Model + finalRole = ollamaResp.Message.Role + + if ollamaResp.Done { + break + } } + return nil + } + + if err := WithRetry(ctx, p.config, attempt); err != nil { + return nil, err } finalContent := completeResponse.String() @@ -246,11 +263,11 @@ func (p *OllamaProvider) CompleteStream(ctx context.Context, req CompletionReque if err != nil { return fmt.Errorf("failed to send request: %w", err) } - defer resp.Body.Close() + defer func() { _ = resp.Body.Close() }() if resp.StatusCode != http.StatusOK { body, _ := io.ReadAll(resp.Body) - return fmt.Errorf("Ollama API error: status %d, body: %s", resp.StatusCode, string(body)) + return fmt.Errorf("%w: ollama API error: status %d, body: %s", ErrProviderUnavailable, resp.StatusCode, string(body)) } decoder := json.NewDecoder(resp.Body) @@ -264,7 +281,7 @@ func (p *OllamaProvider) CompleteStream(ctx context.Context, req CompletionReque } if ollamaResp.Error != "" { - return fmt.Errorf("Ollama API error: %s", ollamaResp.Error) + return fmt.Errorf("%w: ollama API error: %s", ErrProviderRequest, ollamaResp.Error) } // Convert to our format and call callback @@ -293,12 +310,12 @@ func (p *OllamaProvider) IsHealthy(ctx context.Context) error { resp, err := p.client.Do(req) if err != nil { - return fmt.Errorf("Ollama health check failed: %w", err) + return fmt.Errorf("ollama health check failed: %w", err) } - defer resp.Body.Close() + defer func() { _ = resp.Body.Close() }() if resp.StatusCode != http.StatusOK { - return fmt.Errorf("Ollama health check failed: status %d", resp.StatusCode) + return fmt.Errorf("ollama health check failed: status %d", resp.StatusCode) } return nil @@ -692,7 +709,7 @@ func (p *OllamaProvider) PullModel(ctx context.Context, model string) error { if err != nil { return fmt.Errorf("failed to pull model: %w", err) } - defer resp.Body.Close() + defer func() { _ = resp.Body.Close() }() if resp.StatusCode != http.StatusOK { body, _ := io.ReadAll(resp.Body) @@ -746,7 +763,7 @@ func (p *OllamaProvider) DeleteModel(ctx context.Context, model string) error { if err != nil { return fmt.Errorf("failed to delete model: %w", err) } - defer resp.Body.Close() + defer func() { _ = resp.Body.Close() }() if resp.StatusCode != http.StatusOK { body, _ := io.ReadAll(resp.Body) diff --git a/pkg/llm/openai.go b/pkg/llm/openai.go index a31fb77..78cf031 100644 --- a/pkg/llm/openai.go +++ b/pkg/llm/openai.go @@ -10,45 +10,211 @@ import ( "context" "errors" "fmt" + "io" + "net" + "net/http" + "net/url" "strings" + "sync" "time" "github.com/sashabaranov/go-openai" "github.com/sirupsen/logrus" ) +// DefaultOpenAIEndpoint is the public OpenAI API base URL. Any OpenAI-compatible +// deployment (Azure OpenAI, vLLM, a corporate gateway) is reached by setting +// ProviderConfig.Endpoint instead. +const DefaultOpenAIEndpoint = "https://api.openai.com/v1" + +// defaultOpenAITimeout bounds a request when the configuration does not. +const defaultOpenAITimeout = 60 * time.Second + // OpenAIProvider implements the Provider interface for OpenAI type OpenAIProvider struct { - client *openai.Client - config *ProviderConfig - logger *logrus.Logger - models []string - lastSync time.Time + // mu guards everything below it. SetConfig can rebuild the client while + // requests are in flight, and the model cache is shared between callers, + // so both were data races before. + mu sync.RWMutex + client *openai.Client + httpClient *http.Client + config *ProviderConfig + models []string + lastSync time.Time + + logger *logrus.Logger +} + +// openAITransport injects the configured headers on every request and bounds +// the response body. +// +// ProviderConfig.Headers was accepted and then dropped, so the extra headers +// an Azure deployment or a gateway needs never reached the wire. The size cap +// gives this provider the same protection the others get from limitedBody: the +// SDK reads response bodies itself, so the only place to apply it is here. +type openAITransport struct { + base http.RoundTripper + headers map[string]string +} + +func (t *openAITransport) RoundTrip(req *http.Request) (*http.Response, error) { + base := t.base + if base == nil { + base = http.DefaultTransport + } + + // RoundTrip must not mutate the request it is handed. + if len(t.headers) > 0 { + req = req.Clone(req.Context()) + for key, value := range t.headers { + req.Header.Set(key, value) + } + } + + resp, err := base.RoundTrip(req) + if err != nil { + return nil, err + } + resp.Body = cappedBody{Reader: io.LimitReader(resp.Body, MaxResponseBytes), Closer: resp.Body} + return resp, nil +} + +// cappedBody reads at most MaxResponseBytes while still closing the underlying +// body. +type cappedBody struct { + io.Reader + io.Closer } // NewOpenAIProvider creates a new OpenAI provider func NewOpenAIProvider(config *ProviderConfig) (*OpenAIProvider, error) { + if config == nil { + return nil, fmt.Errorf("openai configuration is required") + } if config.APIKey == "" { return nil, fmt.Errorf("OpenAI API key is required") } - clientConfig := openai.DefaultConfig(config.APIKey) - if config.Endpoint != "" { - clientConfig.BaseURL = config.Endpoint + // Record the effective endpoint so GetConfig reports what the client + // actually talks to rather than an empty string. + if config.Endpoint == "" { + config.Endpoint = DefaultOpenAIEndpoint } - client := openai.NewClientWithConfig(clientConfig) - provider := &OpenAIProvider{ - client: client, config: config, logger: logrus.New(), models: []string{}, } + provider.rebuildClient() return provider, nil } +// rebuildClient constructs the SDK client from the current configuration. The +// caller must hold the write lock (or hold no references yet, at construction). +func (p *OpenAIProvider) rebuildClient() { + timeout := p.config.Timeout + if timeout <= 0 { + timeout = defaultOpenAITimeout + } + + headers := make(map[string]string, len(p.config.Headers)) + for key, value := range p.config.Headers { + headers[key] = value + } + + // openai.DefaultConfig installs a bare &http.Client{}, which has no + // timeout: ProviderConfig.Timeout was ignored and an endpoint that + // accepted a connection and then went silent hung the caller forever. + httpClient := &http.Client{ + Timeout: timeout, + Transport: &openAITransport{base: http.DefaultTransport, headers: headers}, + } + + clientConfig := openai.DefaultConfig(p.config.APIKey) + if p.config.Endpoint != "" { + clientConfig.BaseURL = p.config.Endpoint + } + clientConfig.HTTPClient = httpClient + + p.httpClient = httpClient + p.client = openai.NewClientWithConfig(clientConfig) +} + +// state returns the client together with a snapshot of the configuration, so a +// concurrent SetConfig cannot swap them halfway through a request. +func (p *OpenAIProvider) state() (*openai.Client, ProviderConfig) { + p.mu.RLock() + defer p.mu.RUnlock() + return p.client, *p.config +} + +// isReasoningModel reports whether a model belongs to the o-series, which +// accepts a different parameter set from the chat models. +func isReasoningModel(model string) bool { + return strings.HasPrefix(model, "o1") || + strings.HasPrefix(model, "o3") || + strings.HasPrefix(model, "o4") +} + +// classifyOpenAIError maps an SDK failure onto the shared sentinels. +// +// Errors were previously wrapped with fmt.Errorf and nothing else, so a caller +// could not tell a rate limit from a malformed request and every failure was +// equally (un)retryable. The SDK surfaces the HTTP status on *openai.APIError +// and *openai.RequestError; network failures arrive as *url.Error. +func classifyOpenAIError(ctx context.Context, err error) error { + if err == nil { + return nil + } + + // The caller's context ending is never the provider's fault, and retrying + // it cannot help. + if ctxErr := ctx.Err(); ctxErr != nil { + return ctxErr + } + + var apiErr *openai.APIError + if errors.As(err, &apiErr) { + return openAIStatusError(apiErr.HTTPStatusCode, apiErr.Error()) + } + + var reqErr *openai.RequestError + if errors.As(err, &reqErr) { + return openAIStatusError(reqErr.HTTPStatusCode, reqErr.Error()) + } + + // A transport failure (connection refused, DNS, TLS, client timeout) is + // worth another attempt. + var urlErr *url.Error + var netErr net.Error + if errors.As(err, &urlErr) || errors.As(err, &netErr) { + return NewTransportError("OpenAI", err) + } + + // Anything left is local: request validation inside the SDK, or a body + // that did not decode. Neither improves on a retry, so it stays + // unclassified and IsRetryable reports false. + return fmt.Errorf("OpenAI request failed: %w", err) +} + +// openAIStatusError builds a classified error from an HTTP status. The SDK +// discards the response headers, so a provider-supplied Retry-After cannot be +// honored here; WithRetry falls back to its own backoff. +func openAIStatusError(status int, message string) *ProviderError { + kind := classifyStatus(status) + if kind == nil { + kind = ErrProviderUnavailable + } + return &ProviderError{ + Provider: "OpenAI", + StatusCode: status, + Body: message, + kind: kind, + } +} + // GetName returns the provider name func (p *OpenAIProvider) GetName() string { return "openai" @@ -56,47 +222,112 @@ func (p *OpenAIProvider) GetName() string { // GetModels returns available models func (p *OpenAIProvider) GetModels(ctx context.Context) ([]string, error) { + client, cfg := p.state() + // Cache models for 5 minutes - if time.Since(p.lastSync) < 5*time.Minute && len(p.models) > 0 { - return p.models, nil + p.mu.RLock() + cached := append([]string(nil), p.models...) + fresh := time.Since(p.lastSync) < 5*time.Minute + p.mu.RUnlock() + + if fresh && len(cached) > 0 { + return cached, nil } - models, err := p.client.ListModels(ctx) - if err != nil { - return nil, fmt.Errorf("failed to list models: %w", err) + var models []string + attempt := func(ctx context.Context) error { + list, err := client.ListModels(ctx) + if err != nil { + return classifyOpenAIError(ctx, err) + } + models = make([]string, len(list.Models)) + for i, model := range list.Models { + models[i] = model.ID + } + return nil } - p.models = make([]string, len(models.Models)) - for i, model := range models.Models { - p.models[i] = model.ID + if err := WithRetry(ctx, &cfg, attempt); err != nil { + return nil, fmt.Errorf("failed to list models: %w", err) } + p.mu.Lock() + p.models = models p.lastSync = time.Now() - return p.models, nil + p.mu.Unlock() + + // Hand back a copy: the cache must not be mutable through the caller. + return append([]string(nil), models...), nil } // Complete generates a completion func (p *OpenAIProvider) Complete(ctx context.Context, req CompletionRequest) (*CompletionResponse, error) { - openaiReq := p.convertToOpenAIRequest(req) + client, cfg := p.state() - resp, err := p.client.CreateChatCompletion(ctx, openaiReq) + openaiReq, err := p.convertToOpenAIRequest(&cfg, req) if err != nil { - return nil, fmt.Errorf("OpenAI completion failed: %w", err) + return nil, err + } + + // This is the non-streaming path. The SDK rejects a request that carries + // Stream with ErrChatCompletionStreamNotSupported, so passing the flag + // through turned an ordinary completion into a local error; the streaming + // decision belongs to CompleteWithMode. + openaiReq.Stream = false + + p.logger.WithFields(logrus.Fields{ + "endpoint": cfg.Endpoint, + "model": openaiReq.Model, + }).Debug("Sending request to OpenAI") + + var result *CompletionResponse + + // Transient failures (network errors, 5xx, rate limits) are retried + // according to the provider configuration; permanent errors fail fast. + // RetryCount and RetryDelay were accepted and then ignored. + attempt := func(ctx context.Context) error { + resp, err := client.CreateChatCompletion(ctx, openaiReq) + if err != nil { + return classifyOpenAIError(ctx, err) + } + result = p.convertFromOpenAIResponse(resp) + return nil } - return p.convertFromOpenAIResponse(resp), nil + if err := WithRetry(ctx, &cfg, attempt); err != nil { + return nil, err + } + + return result, nil } // CompleteStream generates a streaming completion func (p *OpenAIProvider) CompleteStream(ctx context.Context, req CompletionRequest, callback StreamCallback) error { - openaiReq := p.convertToOpenAIRequest(req) - openaiReq.Stream = true + client, cfg := p.state() - stream, err := p.client.CreateChatCompletionStream(ctx, openaiReq) + openaiReq, err := p.convertToOpenAIRequest(&cfg, req) if err != nil { - return fmt.Errorf("OpenAI streaming failed: %w", err) + return err + } + openaiReq.Stream = true + + // Only opening the stream is retried: once a chunk has reached the + // callback the caller has seen partial output, and replaying the request + // would duplicate it. + var stream *openai.ChatCompletionStream + open := func(ctx context.Context) error { + s, err := client.CreateChatCompletionStream(ctx, openaiReq) + if err != nil { + return classifyOpenAIError(ctx, err) + } + stream = s + return nil } - defer stream.Close() + + if err := WithRetry(ctx, &cfg, open); err != nil { + return err + } + defer func() { _ = stream.Close() }() for { if err := ctx.Err(); err != nil { @@ -104,15 +335,21 @@ func (p *OpenAIProvider) CompleteStream(ctx context.Context, req CompletionReque } response, err := stream.Recv() if err != nil { - if err.Error() == "EOF" || errors.Is(err, context.Canceled) { + // The end of a stream was detected by comparing err.Error() to + // "EOF", which silently swallowed any wrapped error whose text + // happened to match and missed a wrapped io.EOF. + if errors.Is(err, io.EOF) || errors.Is(err, context.Canceled) { break } - return fmt.Errorf("stream error: %w", err) + return classifyOpenAIError(ctx, err) } // Convert to our format and call callback converted := p.convertFromOpenAIStreamResponse(response) if err := callback(converted); err != nil { + // The deferred Close tears down the HTTP body, so abandoning the + // stream here does not leak the connection. + // // Propagate early-exit unwrapped so CollectStream can treat it as success. if IsStreamEarlyExit(err) { return err @@ -126,20 +363,27 @@ func (p *OpenAIProvider) CompleteStream(ctx context.Context, req CompletionReque // IsHealthy checks if the provider is healthy func (p *OpenAIProvider) IsHealthy(ctx context.Context) error { + client, _ := p.state() + // Try to list models as a health check - _, err := p.client.ListModels(ctx) - if err != nil { - return fmt.Errorf("OpenAI health check failed: %w", err) + if _, err := client.ListModels(ctx); err != nil { + return fmt.Errorf("OpenAI health check failed: %w", classifyOpenAIError(ctx, err)) } return nil } // GetConfig returns provider configuration func (p *OpenAIProvider) GetConfig() map[string]interface{} { + p.mu.RLock() + defer p.mu.RUnlock() + return map[string]interface{}{ - "name": p.config.Name, - "type": p.config.Type, - "endpoint": p.config.Endpoint, + "name": p.config.Name, + "type": p.config.Type, + "endpoint": p.config.Endpoint, + // Never hand back the credential: this map is logged and serialized by + // callers. + "api_key": "***masked***", "model": p.config.Model, "temperature": p.config.Temperature, "max_tokens": p.config.MaxTokens, @@ -151,9 +395,25 @@ func (p *OpenAIProvider) GetConfig() map[string]interface{} { // SetConfig updates provider configuration func (p *OpenAIProvider) SetConfig(config map[string]interface{}) error { + p.mu.Lock() + defer p.mu.Unlock() + + // Endpoint, credential and timeout are baked into the SDK client at + // construction, so changing them used to have no effect at all; the client + // is rebuilt below when any of them moves. + rebuild := false + if name, ok := config["name"].(string); ok { p.config.Name = name } + if endpoint, ok := config["endpoint"].(string); ok && endpoint != p.config.Endpoint { + p.config.Endpoint = endpoint + rebuild = true + } + if apiKey, ok := config["api_key"].(string); ok && apiKey != "" && apiKey != p.config.APIKey { + p.config.APIKey = apiKey + rebuild = true + } if model, ok := config["model"].(string); ok { p.config.Model = model } @@ -163,8 +423,9 @@ func (p *OpenAIProvider) SetConfig(config map[string]interface{}) error { if maxTokens, ok := config["max_tokens"].(int); ok { p.config.MaxTokens = maxTokens } - if timeout, ok := config["timeout"].(time.Duration); ok { + if timeout, ok := config["timeout"].(time.Duration); ok && timeout != p.config.Timeout { p.config.Timeout = timeout + rebuild = true } if retryCount, ok := config["retry_count"].(int); ok { p.config.RetryCount = retryCount @@ -172,24 +433,59 @@ func (p *OpenAIProvider) SetConfig(config map[string]interface{}) error { if retryDelay, ok := config["retry_delay"].(time.Duration); ok { p.config.RetryDelay = retryDelay } + if headers, ok := config["headers"].(map[string]string); ok { + p.config.Headers = headers + rebuild = true + } + + if rebuild { + p.rebuildClient() + } return nil } // Close closes the provider and cleans up resources func (p *OpenAIProvider) Close() error { - // OpenAI client doesn't need explicit closing + p.mu.RLock() + defer p.mu.RUnlock() + + // Release keep-alive connections rather than leaving them to the finaliser. + if p.httpClient != nil { + p.httpClient.CloseIdleConnections() + } return nil } // convertToOpenAIRequest converts our request format to OpenAI format -func (p *OpenAIProvider) convertToOpenAIRequest(req CompletionRequest) openai.ChatCompletionRequest { - messages := make([]openai.ChatCompletionMessage, len(req.Messages)) - for i, msg := range req.Messages { - messages[i] = openai.ChatCompletionMessage{ - Role: msg.Role, - Content: msg.Content, - Name: msg.Name, +func (p *OpenAIProvider) convertToOpenAIRequest(cfg *ProviderConfig, req CompletionRequest) (openai.ChatCompletionRequest, error) { + // Use default model if not specified + model := req.Model + if model == "" { + model = cfg.Model + if model == "" { + model = "gpt-3.5-turbo" + } + } + + messages := make([]openai.ChatCompletionMessage, 0, len(req.Messages)+1) + + // CompletionRequest.SystemPrompt was dropped on the floor: a caller that + // set it (the field every other provider honors) had its instructions + // silently discarded. + if req.SystemPrompt != "" { + messages = append(messages, openai.ChatCompletionMessage{ + Role: openai.ChatMessageRoleSystem, + Content: req.SystemPrompt, + }) + } + + for _, msg := range req.Messages { + converted := openai.ChatCompletionMessage{ + Role: msg.Role, + Content: msg.Content, + Name: msg.Name, + ToolCallID: msg.ToolCallID, } // Convert tool calls @@ -205,40 +501,43 @@ func (p *OpenAIProvider) convertToOpenAIRequest(req CompletionRequest) openai.Ch }, } } - messages[i].ToolCalls = toolCalls + converted.ToolCalls = toolCalls } - // Set tool call ID for tool messages - if msg.ToolCallID != "" { - messages[i].ToolCallID = msg.ToolCallID - } + messages = append(messages, converted) } - openaiReq := openai.ChatCompletionRequest{ - Model: req.Model, - Messages: messages, - Temperature: float32(req.Temperature), - MaxTokens: req.MaxTokens, - Stream: req.Stream, - Stop: req.StopSequences, + // An empty message list is a 400 from the API and a wasted round trip. + if len(messages) == 0 { + return openai.ChatCompletionRequest{}, fmt.Errorf("%w: no messages provided", ErrProviderRequest) } - // Use default model if not specified - if openaiReq.Model == "" { - openaiReq.Model = p.config.Model - if openaiReq.Model == "" { - openaiReq.Model = "gpt-3.5-turbo" - } + // Use default temperature and max tokens if not specified + temperature := req.Temperature + if temperature == 0 { + temperature = cfg.Temperature + } + maxTokens := req.MaxTokens + if maxTokens == 0 { + maxTokens = cfg.MaxTokens } - // Use default temperature if not specified - if openaiReq.Temperature == 0 { - openaiReq.Temperature = float32(p.config.Temperature) + openaiReq := openai.ChatCompletionRequest{ + Model: model, + Messages: messages, + Stream: req.Stream, + Stop: req.StopSequences, } - // Use default max tokens if not specified - if openaiReq.MaxTokens == 0 { - openaiReq.MaxTokens = p.config.MaxTokens + if isReasoningModel(model) { + // The o-series rejects max_tokens (max_completion_tokens replaces it) + // and any temperature other than 1. The SDK enforces both locally, so + // sending the chat-model parameter set meant every call to a reasoning + // model failed without ever reaching the API. + openaiReq.MaxCompletionTokens = maxTokens + } else { + openaiReq.MaxTokens = maxTokens + openaiReq.Temperature = float32(temperature) } // Convert tools @@ -261,7 +560,9 @@ func (p *OpenAIProvider) convertToOpenAIRequest(req CompletionRequest) openai.Ch if req.ToolChoice != nil { switch tc := req.ToolChoice.(type) { case string: - if tc == "auto" || tc == "none" { + // Only "auto" and "none" used to survive, so "required" β€” a value + // the API accepts β€” was silently downgraded to the default. + if tc != "" { openaiReq.ToolChoice = tc } case map[string]interface{}: @@ -280,7 +581,7 @@ func (p *OpenAIProvider) convertToOpenAIRequest(req CompletionRequest) openai.Ch } } - return openaiReq + return openaiReq, nil } // convertFromOpenAIResponse converts OpenAI response to our format @@ -368,7 +669,7 @@ func (p *OpenAIProvider) convertFromOpenAIStreamResponse(resp openai.ChatComplet } } - return CompletionResponse{ + converted := CompletionResponse{ ID: resp.ID, Object: resp.Object, Created: resp.Created, @@ -376,6 +677,18 @@ func (p *OpenAIProvider) convertFromOpenAIStreamResponse(resp openai.ChatComplet Choices: choices, SystemFingerprint: resp.SystemFingerprint, } + + // Usage arrives on the final chunk when stream_options.include_usage is + // set; it used to be discarded, leaving callers with no token counts. + if resp.Usage != nil { + converted.Usage = Usage{ + PromptTokens: resp.Usage.PromptTokens, + CompletionTokens: resp.Usage.CompletionTokens, + TotalTokens: resp.Usage.TotalTokens, + } + } + + return converted } // GetDefaultModels returns commonly used OpenAI models @@ -399,6 +712,9 @@ func (p *OpenAIProvider) SupportsStreaming() bool { // GetStreamingConfig returns the current streaming configuration func (p *OpenAIProvider) GetStreamingConfig() *StreamingConfig { + p.mu.Lock() + defer p.mu.Unlock() + if p.config.Streaming == nil { p.config.Streaming = DefaultStreamingConfig() } @@ -410,6 +726,10 @@ func (p *OpenAIProvider) SetStreamingConfig(config *StreamingConfig) error { if config == nil { return fmt.Errorf("streaming config cannot be nil") } + + p.mu.Lock() + defer p.mu.Unlock() + p.config.Streaming = config return nil } @@ -459,7 +779,7 @@ func (p *OpenAIProvider) completeNonStreaming(ctx context.Context, req Completio } // completeStreamingCollected forces streaming but collects all chunks into single response. -// When req.EarlyExit fires, the remainder of the token stream is cancelled and the +// When req.EarlyExit fires, the remainder of the token stream is canceled and the // accumulated content/tool-calls are returned successfully (FinishReason=early_exit). func (p *OpenAIProvider) completeStreamingCollected(ctx context.Context, req CompletionRequest) (*CompletionResponse, error) { return CollectStream(ctx, p.CompleteStream, req) @@ -519,7 +839,7 @@ func (p *OpenAIProvider) ValidateModel(model string) error { } // Check if it's a known OpenAI model pattern - validPrefixes := []string{"gpt-", "text-", "code-", "davinci", "curie", "babbage", "ada"} + validPrefixes := []string{"gpt-", "text-", "code-", "davinci", "curie", "babbage", "ada", "o1", "o3", "o4", "chatgpt-"} for _, prefix := range validPrefixes { if strings.HasPrefix(model, prefix) { return nil diff --git a/pkg/llm/openai_test.go b/pkg/llm/openai_test.go new file mode 100644 index 0000000..9410454 --- /dev/null +++ b/pkg/llm/openai_test.go @@ -0,0 +1,965 @@ +// Copyright (c) 2024 GoLangGraph Team +// +// Licensed under the MIT License. See LICENSE file in the project root for full license information. + +package llm + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" + "net/http/httptest" + "strings" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/sashabaranov/go-openai" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// openAIAPI is a stand-in for the OpenAI API. The provider under test is the +// real one: it builds real requests through the SDK, speaks the real wire +// format and parses real responses. +type openAIAPI struct { + server *httptest.Server + requests atomic.Int32 + + // Recorded details of the most recent request. The handler runs on the + // server's goroutine, so access is guarded. + mu sync.Mutex + lastPath string + lastMethod string + lastHeader http.Header + lastBody map[string]interface{} +} + +func newOpenAIAPI(t *testing.T, handler http.HandlerFunc) *openAIAPI { + t.Helper() + api := &openAIAPI{} + api.server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + api.requests.Add(1) + + var decoded map[string]interface{} + if r.Body != nil { + if raw, err := io.ReadAll(r.Body); err == nil && len(raw) > 0 { + _ = json.Unmarshal(raw, &decoded) + } + } + + api.mu.Lock() + api.lastPath = r.URL.Path + api.lastMethod = r.Method + api.lastHeader = r.Header.Clone() + if decoded != nil { + api.lastBody = decoded + } + api.mu.Unlock() + + handler(w, r) + })) + t.Cleanup(api.server.Close) + return api +} + +func (api *openAIAPI) body(t *testing.T) map[string]interface{} { + t.Helper() + api.mu.Lock() + defer api.mu.Unlock() + require.NotNil(t, api.lastBody, "no request body was recorded") + return api.lastBody +} + +func (api *openAIAPI) header(t *testing.T, name string) string { + t.Helper() + api.mu.Lock() + defer api.mu.Unlock() + require.NotNil(t, api.lastHeader, "no request was recorded") + return api.lastHeader.Get(name) +} + +func (api *openAIAPI) path(t *testing.T) string { + t.Helper() + api.mu.Lock() + defer api.mu.Unlock() + return api.lastPath +} + +// provider builds an OpenAI provider pointed at the fake API. Pointing the SDK +// client at a test server is only possible because the provider honors +// ProviderConfig.Endpoint. +func (api *openAIAPI) provider(t *testing.T, mutate func(*ProviderConfig)) *OpenAIProvider { + t.Helper() + cfg := DefaultProviderConfig() + cfg.Type = "openai" + cfg.APIKey = "test-key" // pragma: allowlist secret + cfg.Model = "gpt-4o-mini" + cfg.Endpoint = api.server.URL + cfg.RetryCount = 0 + cfg.RetryDelay = time.Millisecond + cfg.Timeout = 5 * time.Second + if mutate != nil { + mutate(cfg) + } + p, err := NewOpenAIProvider(cfg) + require.NoError(t, err) + t.Cleanup(func() { _ = p.Close() }) + return p +} + +// openAIChatSuccess answers with a well-formed chat completion. +func openAIChatSuccess(content string) http.HandlerFunc { + return func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(map[string]interface{}{ + "id": "chatcmpl-test", + "object": "chat.completion", + "created": 1700000000, + "model": "gpt-4o-mini", + "choices": []map[string]interface{}{{ + "index": 0, + "message": map[string]interface{}{"role": "assistant", "content": content}, + "finish_reason": "stop", + }}, + "usage": map[string]int{ + "prompt_tokens": 7, "completion_tokens": 11, "total_tokens": 18, + }, + "system_fingerprint": "fp_test", + }) + } +} + +// openAIError answers with the API's error envelope, which the SDK decodes +// into *openai.APIError. +func openAIError(status int, message string) http.HandlerFunc { + return func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(status) + _ = json.NewEncoder(w).Encode(map[string]interface{}{ + "error": map[string]interface{}{ + "message": message, + "type": "test_error", + "code": "test_code", + }, + }) + } +} + +// openAIStream answers with server-sent events, one content delta per chunk. +func openAIStream(chunks ...string) http.HandlerFunc { + return func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + flusher, _ := w.(http.Flusher) + for i, chunk := range chunks { + delta := map[string]interface{}{"content": chunk} + if i == 0 { + delta["role"] = "assistant" + } + frame, _ := json.Marshal(map[string]interface{}{ + "id": "chatcmpl-stream", + "object": "chat.completion.chunk", + "created": 1700000000, + "model": "gpt-4o-mini", + "choices": []map[string]interface{}{{"index": 0, "delta": delta}}, + }) + _, _ = fmt.Fprintf(w, "data: %s\n\n", frame) + if flusher != nil { + flusher.Flush() + } + } + _, _ = fmt.Fprint(w, "data: [DONE]\n\n") + if flusher != nil { + flusher.Flush() + } + } +} + +// openAIRoutes dispatches between the chat and models endpoints. +func openAIRoutes(chat, models http.HandlerFunc) http.HandlerFunc { + return func(w http.ResponseWriter, r *http.Request) { + switch { + case strings.HasSuffix(r.URL.Path, "/chat/completions"): + chat(w, r) + case strings.HasSuffix(r.URL.Path, "/models"): + models(w, r) + default: + http.NotFound(w, r) + } + } +} + +func openAITestRequest() CompletionRequest { + return CompletionRequest{ + Messages: []Message{{Role: "user", Content: "Hello, how are you?"}}, + } +} + +// The completion must come from the API, proving the provider actually calls it. +func TestOpenAI_CallsTheAPI(t *testing.T) { + api := newOpenAIAPI(t, openAIChatSuccess("a genuine model reply")) + p := api.provider(t, nil) + + resp, err := p.Complete(context.Background(), openAITestRequest()) + require.NoError(t, err) + + assert.EqualValues(t, 1, api.requests.Load(), "the provider must actually call the API") + assert.Equal(t, "/chat/completions", api.path(t)) + require.NotEmpty(t, resp.Choices) + assert.Equal(t, "a genuine model reply", resp.Choices[0].Message.Content) + assert.Equal(t, "assistant", resp.Choices[0].Message.Role) + assert.Equal(t, "stop", resp.Choices[0].FinishReason) + assert.Equal(t, "chatcmpl-test", resp.ID) + assert.Equal(t, "fp_test", resp.SystemFingerprint) + assert.Equal(t, 7, resp.Usage.PromptTokens) + assert.Equal(t, 11, resp.Usage.CompletionTokens) + assert.Equal(t, 18, resp.Usage.TotalTokens) +} + +// Model, messages and generation settings must reach the API, along with the +// credential. +func TestOpenAI_RequestShape(t *testing.T) { + api := newOpenAIAPI(t, openAIChatSuccess("ok")) + p := api.provider(t, nil) + + _, err := p.Complete(context.Background(), CompletionRequest{ + Model: "gpt-4o", + Messages: []Message{{Role: "user", Content: "hi"}, {Role: "assistant", Content: "hello"}}, + Temperature: 0.25, + MaxTokens: 512, + StopSequences: []string{"END"}, + }) + require.NoError(t, err) + + body := api.body(t) + assert.Equal(t, "gpt-4o", body["model"]) + assert.InDelta(t, 0.25, body["temperature"], 1e-6) + assert.EqualValues(t, 512, body["max_tokens"]) + assert.Contains(t, fmt.Sprint(body["stop"]), "END") + + messages, ok := body["messages"].([]interface{}) + require.True(t, ok, "messages must be sent: %v", body) + require.Len(t, messages, 2) + first := messages[0].(map[string]interface{}) + assert.Equal(t, "user", first["role"]) + assert.Equal(t, "hi", first["content"]) + assert.Equal(t, "assistant", messages[1].(map[string]interface{})["role"]) + + assert.Equal(t, "Bearer test-key", api.header(t, "Authorization")) +} + +// The configured model is used when the request does not name one. +func TestOpenAI_FallsBackToConfiguredModel(t *testing.T) { + api := newOpenAIAPI(t, openAIChatSuccess("ok")) + p := api.provider(t, func(c *ProviderConfig) { + c.Model = "gpt-4-turbo" + c.MaxTokens = 321 + c.Temperature = 0.9 + }) + + _, err := p.Complete(context.Background(), openAITestRequest()) + require.NoError(t, err) + + body := api.body(t) + assert.Equal(t, "gpt-4-turbo", body["model"]) + assert.EqualValues(t, 321, body["max_tokens"]) + assert.InDelta(t, 0.9, body["temperature"], 1e-6) +} + +// CompletionRequest.SystemPrompt is part of the shared request type and was +// dropped entirely by this provider, so the caller's instructions never +// reached the model. +func TestOpenAI_SystemPromptIsSent(t *testing.T) { + api := newOpenAIAPI(t, openAIChatSuccess("ok")) + p := api.provider(t, nil) + + _, err := p.Complete(context.Background(), CompletionRequest{ + SystemPrompt: "You are terse.", + Messages: []Message{{Role: "user", Content: "hi"}}, + }) + require.NoError(t, err) + + messages, ok := api.body(t)["messages"].([]interface{}) + require.True(t, ok) + require.Len(t, messages, 2, "the system prompt must be prepended as a message") + first := messages[0].(map[string]interface{}) + assert.Equal(t, "system", first["role"]) + assert.Equal(t, "You are terse.", first["content"]) +} + +// Reasoning models reject max_tokens and a non-default temperature; the SDK +// enforces that locally, so sending the chat-model parameter set meant every +// o-series call failed without ever reaching the API. +func TestOpenAI_ReasoningModelParameters(t *testing.T) { + api := newOpenAIAPI(t, openAIChatSuccess("thought about it")) + p := api.provider(t, nil) + + resp, err := p.Complete(context.Background(), CompletionRequest{ + Model: "o3-mini", + Messages: []Message{{Role: "user", Content: "hi"}}, + MaxTokens: 256, + Temperature: 0.7, + }) + require.NoError(t, err, "a reasoning model request must reach the API") + assert.Equal(t, "thought about it", resp.Choices[0].Message.Content) + + body := api.body(t) + assert.EqualValues(t, 256, body["max_completion_tokens"]) + assert.NotContains(t, body, "max_tokens", "max_tokens is rejected by reasoning models") + assert.NotContains(t, body, "temperature", "reasoning models only accept the default temperature") +} + +// Complete is the non-streaming path. The SDK refuses a request that carries +// the stream flag, so passing it through turned an ordinary completion into a +// local error before any request was made. +func TestOpenAI_CompleteIgnoresStreamFlag(t *testing.T) { + api := newOpenAIAPI(t, openAIChatSuccess("ok")) + p := api.provider(t, nil) + + req := openAITestRequest() + req.Stream = true + + resp, err := p.Complete(context.Background(), req) + require.NoError(t, err) + assert.Equal(t, "ok", resp.Choices[0].Message.Content) + assert.EqualValues(t, 1, api.requests.Load()) + + assert.NotContains(t, api.body(t), "stream", "Complete must not ask for a stream") +} + +// Tools go out on the wire and tool calls come back parsed. +func TestOpenAI_ToolCalls(t *testing.T) { + api := newOpenAIAPI(t, func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(map[string]interface{}{ + "id": "chatcmpl-tools", + "object": "chat.completion", + "model": "gpt-4o-mini", + "choices": []map[string]interface{}{{ + "index": 0, + "message": map[string]interface{}{ + "role": "assistant", + "tool_calls": []map[string]interface{}{{ + "id": "call_1", + "type": "function", + "function": map[string]string{"name": "get_weather", "arguments": `{"city":"Paris"}`}, + }}, + }, + "finish_reason": "tool_calls", + }}, + }) + }) + p := api.provider(t, nil) + + resp, err := p.Complete(context.Background(), CompletionRequest{ + Messages: []Message{{Role: "user", Content: "weather?"}}, + Tools: []ToolDefinition{{ + Type: "function", + Function: Function{ + Name: "get_weather", + Description: "look up the weather", + Parameters: map[string]interface{}{"type": "object"}, + }, + }}, + // "required" is a value the API accepts; it used to be discarded + // because only "auto" and "none" were passed through. + ToolChoice: "required", + }) + require.NoError(t, err) + + body := api.body(t) + tools, ok := body["tools"].([]interface{}) + require.True(t, ok, "tools must be sent: %v", body) + require.Len(t, tools, 1) + assert.Equal(t, "required", body["tool_choice"]) + + require.NotEmpty(t, resp.Choices) + require.Len(t, resp.Choices[0].Message.ToolCalls, 1) + call := resp.Choices[0].Message.ToolCalls[0] + assert.Equal(t, "call_1", call.ID) + assert.Equal(t, "function", call.Type) + assert.Equal(t, "get_weather", call.Function.Name) + assert.JSONEq(t, `{"city":"Paris"}`, call.Function.Arguments) + assert.Equal(t, "tool_calls", resp.Choices[0].FinishReason) +} + +// Configured headers reach the wire; they used to be accepted and dropped, +// leaving no way to address a gateway that needs them. +func TestOpenAI_CustomHeaders(t *testing.T) { + api := newOpenAIAPI(t, openAIChatSuccess("ok")) + p := api.provider(t, func(c *ProviderConfig) { + c.Headers = map[string]string{"X-Gateway-Tenant": "acme"} + }) + + _, err := p.Complete(context.Background(), openAITestRequest()) + require.NoError(t, err) + assert.Equal(t, "acme", api.header(t, "X-Gateway-Tenant")) +} + +// Errors must be classified so callers can tell a rate limit from a bad +// request instead of matching on message text. +func TestOpenAI_ErrorClassification(t *testing.T) { + cases := []struct { + status int + want error + }{ + {http.StatusBadRequest, ErrProviderRequest}, + {http.StatusUnauthorized, ErrProviderAuth}, + {http.StatusForbidden, ErrProviderAuth}, + {http.StatusTooManyRequests, ErrRateLimited}, + {http.StatusInternalServerError, ErrProviderUnavailable}, + {http.StatusServiceUnavailable, ErrProviderUnavailable}, + } + + for _, tc := range cases { + t.Run(fmt.Sprint(tc.status), func(t *testing.T) { + api := newOpenAIAPI(t, openAIError(tc.status, "upstream said no")) + p := api.provider(t, nil) + + _, err := p.Complete(context.Background(), openAITestRequest()) + require.Error(t, err) + assert.True(t, errors.Is(err, tc.want), "status %d: got %v", tc.status, err) + assert.Contains(t, err.Error(), "upstream said no", "the provider's message must survive") + + var pe *ProviderError + require.True(t, errors.As(err, &pe), "the HTTP status must be recoverable") + assert.Equal(t, tc.status, pe.StatusCode) + }) + } +} + +// A non-JSON error body arrives as a different SDK error type; it must still +// be classified by status. +func TestOpenAI_ErrorClassificationWithoutJSONBody(t *testing.T) { + api := newOpenAIAPI(t, func(w http.ResponseWriter, r *http.Request) { + http.Error(w, "502 Bad Gateway from the proxy", http.StatusBadGateway) + }) + p := api.provider(t, nil) + + _, err := p.Complete(context.Background(), openAITestRequest()) + require.Error(t, err) + assert.True(t, errors.Is(err, ErrProviderUnavailable), "got %v", err) + assert.True(t, IsRetryable(err)) +} + +// Transient failures are retried according to the configuration; RetryCount +// and RetryDelay used to be accepted and ignored. +func TestOpenAI_RetriesTransientFailures(t *testing.T) { + var calls atomic.Int32 + api := newOpenAIAPI(t, func(w http.ResponseWriter, r *http.Request) { + if calls.Add(1) < 3 { + openAIError(http.StatusServiceUnavailable, "try later")(w, r) + return + } + openAIChatSuccess("recovered")(w, r) + }) + p := api.provider(t, func(c *ProviderConfig) { c.RetryCount = 3 }) + + resp, err := p.Complete(context.Background(), openAITestRequest()) + require.NoError(t, err) + assert.Equal(t, "recovered", resp.Choices[0].Message.Content) + assert.EqualValues(t, 3, api.requests.Load(), "the failed attempts must be retried") +} + +// A rate limit is transient too. +func TestOpenAI_RetriesRateLimit(t *testing.T) { + var calls atomic.Int32 + api := newOpenAIAPI(t, func(w http.ResponseWriter, r *http.Request) { + if calls.Add(1) == 1 { + openAIError(http.StatusTooManyRequests, "slow down")(w, r) + return + } + openAIChatSuccess("second time lucky")(w, r) + }) + p := api.provider(t, func(c *ProviderConfig) { c.RetryCount = 2 }) + + resp, err := p.Complete(context.Background(), openAITestRequest()) + require.NoError(t, err) + assert.Equal(t, "second time lucky", resp.Choices[0].Message.Content) + assert.EqualValues(t, 2, api.requests.Load()) +} + +// A permanent failure must fail fast: retrying a malformed request only burns +// the budget and delays the error. +func TestOpenAI_DoesNotRetryPermanentFailures(t *testing.T) { + for _, status := range []int{http.StatusBadRequest, http.StatusUnauthorized} { + t.Run(fmt.Sprint(status), func(t *testing.T) { + api := newOpenAIAPI(t, openAIError(status, "nope")) + p := api.provider(t, func(c *ProviderConfig) { c.RetryCount = 5 }) + + _, err := p.Complete(context.Background(), openAITestRequest()) + require.Error(t, err) + assert.False(t, IsRetryable(err)) + assert.EqualValues(t, 1, api.requests.Load(), "a permanent failure must not be retried") + }) + } +} + +// The retry budget is finite: once it is spent the classified error surfaces. +func TestOpenAI_RetryBudgetExhausted(t *testing.T) { + api := newOpenAIAPI(t, openAIError(http.StatusServiceUnavailable, "still down")) + p := api.provider(t, func(c *ProviderConfig) { c.RetryCount = 2 }) + + _, err := p.Complete(context.Background(), openAITestRequest()) + require.Error(t, err) + assert.True(t, errors.Is(err, ErrProviderUnavailable)) + assert.EqualValues(t, 3, api.requests.Load(), "one initial attempt plus two retries") +} + +// Cancelling the caller's context must stop the call promptly and must not be +// retried. +func TestOpenAI_Cancellation(t *testing.T) { + release := make(chan struct{}) + api := newOpenAIAPI(t, func(w http.ResponseWriter, r *http.Request) { + <-release + }) + defer close(release) + + p := api.provider(t, func(c *ProviderConfig) { + c.Timeout = 30 * time.Second + c.RetryCount = 3 + c.RetryDelay = time.Second + }) + + ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond) + defer cancel() + + start := time.Now() + _, err := p.Complete(ctx, openAITestRequest()) + require.Error(t, err) + assert.True(t, errors.Is(err, context.DeadlineExceeded), "got %v", err) + assert.Less(t, time.Since(start), 5*time.Second, "cancellation must not wait out the retry budget") + assert.EqualValues(t, 1, api.requests.Load(), "a canceled call must not be retried") +} + +// ProviderConfig.Timeout must bound the request. The SDK installs a client +// with no timeout of its own, so an endpoint that accepted the connection and +// then went silent hung the caller indefinitely. +func TestOpenAI_TimeoutIsApplied(t *testing.T) { + api := newOpenAIAPI(t, func(w http.ResponseWriter, r *http.Request) { + select { + case <-r.Context().Done(): + case <-time.After(3 * time.Second): + } + }) + p := api.provider(t, func(c *ProviderConfig) { c.Timeout = 100 * time.Millisecond }) + + start := time.Now() + _, err := p.Complete(context.Background(), openAITestRequest()) + elapsed := time.Since(start) + + require.Error(t, err) + assert.Less(t, elapsed, 2*time.Second, "the configured timeout must bound the request") + assert.True(t, errors.Is(err, ErrProviderUnavailable), "a timeout is transient: %v", err) +} + +// A body that does not decode is a real error, and retrying it will not help. +func TestOpenAI_MalformedResponse(t *testing.T) { + api := newOpenAIAPI(t, func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte("{not json")) + }) + p := api.provider(t, func(c *ProviderConfig) { c.RetryCount = 3 }) + + _, err := p.Complete(context.Background(), openAITestRequest()) + require.Error(t, err) + assert.Contains(t, err.Error(), "OpenAI") + assert.Contains(t, err.Error(), "invalid character") + assert.False(t, IsRetryable(err)) + assert.EqualValues(t, 1, api.requests.Load()) +} + +// An empty request is rejected locally rather than spending a round trip. +func TestOpenAI_EmptyMessagesRejected(t *testing.T) { + api := newOpenAIAPI(t, openAIChatSuccess("ok")) + p := api.provider(t, nil) + + _, err := p.Complete(context.Background(), CompletionRequest{}) + require.Error(t, err) + assert.True(t, errors.Is(err, ErrProviderRequest)) + assert.Zero(t, api.requests.Load(), "an invalid request must not reach the API") +} + +// Streaming must deliver every chunk in order. +func TestOpenAI_Streaming(t *testing.T) { + api := newOpenAIAPI(t, openAIStream("Hello", " there", "!")) + p := api.provider(t, nil) + + var chunks []string + err := p.CompleteStream(context.Background(), openAITestRequest(), func(chunk CompletionResponse) error { + require.NotEmpty(t, chunk.Choices) + chunks = append(chunks, chunk.Choices[0].Delta.Content) + return nil + }) + require.NoError(t, err) + + assert.Equal(t, []string{"Hello", " there", "!"}, chunks) + assert.Equal(t, "Hello there!", strings.Join(chunks, "")) + assert.Equal(t, "/chat/completions", api.path(t)) + assert.Equal(t, true, api.body(t)["stream"], "a streaming call must ask for a stream") +} + +// Opening a stream is retried like any other transient failure; no chunk has +// been delivered yet, so replaying the request is safe. +func TestOpenAI_StreamingRetriesOnOpen(t *testing.T) { + var calls atomic.Int32 + api := newOpenAIAPI(t, func(w http.ResponseWriter, r *http.Request) { + if calls.Add(1) == 1 { + openAIError(http.StatusServiceUnavailable, "warming up")(w, r) + return + } + openAIStream("late", " but", " here")(w, r) + }) + p := api.provider(t, func(c *ProviderConfig) { c.RetryCount = 2 }) + + var got strings.Builder + err := p.CompleteStream(context.Background(), openAITestRequest(), func(chunk CompletionResponse) error { + got.WriteString(chunk.Choices[0].Delta.Content) + return nil + }) + require.NoError(t, err) + assert.Equal(t, "late but here", got.String()) + assert.EqualValues(t, 2, api.requests.Load()) +} + +// A rejected streaming request must surface a classified error, not a bare +// wrapped string. +func TestOpenAI_StreamingErrorStatus(t *testing.T) { + api := newOpenAIAPI(t, openAIError(http.StatusUnauthorized, "bad key")) + p := api.provider(t, func(c *ProviderConfig) { c.RetryCount = 2 }) + + called := 0 + err := p.CompleteStream(context.Background(), openAITestRequest(), func(chunk CompletionResponse) error { + called++ + return nil + }) + require.Error(t, err) + assert.True(t, errors.Is(err, ErrProviderAuth), "got %v", err) + assert.Zero(t, called) + assert.EqualValues(t, 1, api.requests.Load(), "an auth failure must not be retried") +} + +// A callback that fails aborts the stream, surfaces its error, and closes the +// HTTP body β€” the server sees the request context end. +func TestOpenAI_StreamingCallbackError(t *testing.T) { + serverDone := make(chan struct{}) + abandoned := make(chan struct{}) + var once sync.Once + + api := newOpenAIAPI(t, func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + flusher, _ := w.(http.Flusher) + frame := `{"id":"chatcmpl-stream","object":"chat.completion.chunk","choices":[{"index":0,"delta":{"content":"x"}}]}` + _, _ = fmt.Fprintf(w, "data: %s\n\n", frame) + if flusher != nil { + flusher.Flush() + } + + // Hold the response open: the client abandoning the stream is what + // ends the request context. + select { + case <-r.Context().Done(): + once.Do(func() { close(abandoned) }) + case <-serverDone: + } + }) + defer close(serverDone) + + p := api.provider(t, nil) + + sentinel := errors.New("consumer gave up") + calls := 0 + err := p.CompleteStream(context.Background(), openAITestRequest(), func(chunk CompletionResponse) error { + calls++ + return sentinel + }) + require.Error(t, err) + assert.ErrorIs(t, err, sentinel) + assert.Equal(t, 1, calls, "the stream must stop at the first callback error") + + select { + case <-abandoned: + case <-time.After(5 * time.Second): + t.Fatal("the stream was not closed after the callback failed") + } +} + +// The API reports a mid-generation failure as an error frame inside an +// otherwise successful stream. That must surface as a classified error rather +// than a silently truncated answer. +func TestOpenAI_StreamingInBandError(t *testing.T) { + api := newOpenAIAPI(t, func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + flusher, _ := w.(http.Flusher) + _, _ = fmt.Fprint(w, `data: {"id":"c","object":"chat.completion.chunk","choices":[{"index":0,"delta":{"content":"partial"}}]}`+"\n\n") + if flusher != nil { + flusher.Flush() + } + _, _ = fmt.Fprint(w, `data: {"error":{"message":"the server had a problem","type":"server_error","code":null}}`+"\n\n") + if flusher != nil { + flusher.Flush() + } + }) + p := api.provider(t, nil) + + var delivered []string + err := p.CompleteStream(context.Background(), openAITestRequest(), func(chunk CompletionResponse) error { + delivered = append(delivered, chunk.Choices[0].Delta.Content) + return nil + }) + require.Error(t, err, "a failed generation must not look like a complete one") + assert.Equal(t, []string{"partial"}, delivered) + assert.Contains(t, err.Error(), "the server had a problem") + + var pe *ProviderError + assert.True(t, errors.As(err, &pe), "an in-band stream error must still be classified: %v", err) +} + +// Streaming collected into a single response must assemble the whole text and +// carry a usable assistant message. +func TestOpenAI_CompleteWithModeForcedStreaming(t *testing.T) { + api := newOpenAIAPI(t, openAIStream("one ", "two ", "three")) + p := api.provider(t, nil) + + resp, err := p.CompleteWithMode(context.Background(), openAITestRequest(), StreamModeForced) + require.NoError(t, err) + require.NotNil(t, resp) + require.NotEmpty(t, resp.Choices) + assert.Equal(t, "one two three", resp.Choices[0].Message.Content) + assert.Equal(t, "assistant", resp.Choices[0].Message.Role) + assert.Equal(t, "chat.completion", resp.Object) + assert.Empty(t, resp.Choices[0].Delta.Content) +} + +// Health must reflect the API rather than a hardcoded success. +func TestOpenAI_Health(t *testing.T) { + t.Run("healthy", func(t *testing.T) { + api := newOpenAIAPI(t, openAIRoutes(openAIChatSuccess("ok"), func(w http.ResponseWriter, r *http.Request) { + _ = json.NewEncoder(w).Encode(map[string]interface{}{ + "object": "list", + "data": []map[string]interface{}{{"id": "gpt-4o-mini", "object": "model"}}, + }) + })) + p := api.provider(t, nil) + + require.NoError(t, p.IsHealthy(context.Background())) + assert.EqualValues(t, 1, api.requests.Load(), "the health check must call the API") + assert.Equal(t, "/models", api.path(t)) + }) + + t.Run("rejected credentials", func(t *testing.T) { + api := newOpenAIAPI(t, openAIError(http.StatusUnauthorized, "bad key")) + p := api.provider(t, nil) + + err := p.IsHealthy(context.Background()) + require.Error(t, err, "an unhealthy provider must not report healthy") + assert.True(t, errors.Is(err, ErrProviderAuth), "got %v", err) + }) + + t.Run("unreachable", func(t *testing.T) { + api := newOpenAIAPI(t, openAIChatSuccess("ok")) + p := api.provider(t, nil) + api.server.Close() // the endpoint stops answering + + err := p.IsHealthy(context.Background()) + require.Error(t, err) + assert.True(t, errors.Is(err, ErrProviderUnavailable), "got %v", err) + }) +} + +// GetModels reports what the API returns, and caches it. +func TestOpenAI_GetModels(t *testing.T) { + api := newOpenAIAPI(t, openAIRoutes(openAIChatSuccess("ok"), func(w http.ResponseWriter, r *http.Request) { + _ = json.NewEncoder(w).Encode(map[string]interface{}{ + "object": "list", + "data": []map[string]interface{}{ + {"id": "gpt-4o", "object": "model"}, + {"id": "o3-mini", "object": "model"}, + }, + }) + })) + p := api.provider(t, nil) + + models, err := p.GetModels(context.Background()) + require.NoError(t, err) + assert.Equal(t, []string{"gpt-4o", "o3-mini"}, models) + + // Mutating the returned slice must not corrupt the cache. + models[0] = "tampered" + + again, err := p.GetModels(context.Background()) + require.NoError(t, err) + assert.Equal(t, []string{"gpt-4o", "o3-mini"}, again) + assert.EqualValues(t, 1, api.requests.Load(), "the model list must be cached") +} + +func TestOpenAI_GetModelsErrorIsClassified(t *testing.T) { + api := newOpenAIAPI(t, openAIError(http.StatusUnauthorized, "bad key")) + p := api.provider(t, nil) + + _, err := p.GetModels(context.Background()) + require.Error(t, err) + assert.True(t, errors.Is(err, ErrProviderAuth), "got %v", err) +} + +// The endpoint must be honored, and reported, so a gateway or Azure +// deployment can be addressed at all. +func TestOpenAI_EndpointIsHonored(t *testing.T) { + api := newOpenAIAPI(t, openAIChatSuccess("from the gateway")) + p := api.provider(t, nil) + + resp, err := p.Complete(context.Background(), openAITestRequest()) + require.NoError(t, err) + assert.Equal(t, "from the gateway", resp.Choices[0].Message.Content) + assert.Equal(t, api.server.URL, p.GetConfig()["endpoint"]) +} + +// Without an endpoint the provider must still report where it points. +func TestOpenAI_DefaultEndpointReported(t *testing.T) { + p, err := NewOpenAIProvider(&ProviderConfig{Type: "openai", APIKey: "test-key"}) // pragma: allowlist secret + require.NoError(t, err) + t.Cleanup(func() { _ = p.Close() }) + + assert.Equal(t, DefaultOpenAIEndpoint, p.GetConfig()["endpoint"]) +} + +// Re-pointing the provider must take effect: the endpoint and credential are +// baked into the SDK client, so updating them without rebuilding it silently +// did nothing. +func TestOpenAI_SetConfigRepointsClient(t *testing.T) { + first := newOpenAIAPI(t, openAIChatSuccess("from the first server")) + second := newOpenAIAPI(t, openAIChatSuccess("from the second server")) + p := first.provider(t, nil) + + resp, err := p.Complete(context.Background(), openAITestRequest()) + require.NoError(t, err) + assert.Equal(t, "from the first server", resp.Choices[0].Message.Content) + + require.NoError(t, p.SetConfig(map[string]interface{}{ + "endpoint": second.server.URL, + "api_key": "rotated-key", // pragma: allowlist secret + })) + + resp, err = p.Complete(context.Background(), openAITestRequest()) + require.NoError(t, err) + assert.Equal(t, "from the second server", resp.Choices[0].Message.Content) + assert.Equal(t, "Bearer rotated-key", second.header(t, "Authorization")) + assert.EqualValues(t, 1, first.requests.Load(), "no further traffic to the old endpoint") +} + +// Credentials must never appear in the configuration the API exposes. +func TestOpenAI_ConfigMasksAPIKey(t *testing.T) { + api := newOpenAIAPI(t, openAIChatSuccess("ok")) + p := api.provider(t, nil) + + cfg := p.GetConfig() + require.Contains(t, cfg, "api_key", "the credential must be reported as masked, not omitted") + assert.Equal(t, "***masked***", cfg["api_key"]) + assert.NotContains(t, fmt.Sprint(cfg), "test-key") +} + +func TestOpenAI_RequiresAPIKey(t *testing.T) { + _, err := NewOpenAIProvider(&ProviderConfig{Type: "openai", Model: "gpt-4o"}) + assert.Error(t, err, "an empty API key must be rejected at construction") + + _, err = NewOpenAIProvider(nil) + assert.Error(t, err, "a nil configuration must not panic") +} + +// o-series models are valid OpenAI models; they used to be rejected by the +// prefix list. +func TestOpenAI_ValidateModel(t *testing.T) { + p, err := NewOpenAIProvider(&ProviderConfig{Type: "openai", APIKey: "test-key"}) // pragma: allowlist secret + require.NoError(t, err) + t.Cleanup(func() { _ = p.Close() }) + + for _, model := range []string{"gpt-4o", "o1-preview", "o3-mini", "o4-mini"} { + assert.NoError(t, p.ValidateModel(model), "model %s", model) + } + assert.Error(t, p.ValidateModel("")) + assert.Error(t, p.ValidateModel("llama2")) +} + +// The three defects above are all consequences of SDK behavior that is easy +// to get wrong. Pinning it here documents why the provider translates the +// request the way it does, and tells us when a workaround can be dropped. +func TestOpenAI_SDKConstraintsThisProviderWorksAround(t *testing.T) { + api := newOpenAIAPI(t, openAIChatSuccess("ok")) + client := openai.NewClientWithConfig(func() openai.ClientConfig { + c := openai.DefaultConfig("test-key") // pragma: allowlist secret + c.BaseURL = api.server.URL + return c + }()) + + t.Run("max_tokens is rejected for reasoning models", func(t *testing.T) { + _, err := client.CreateChatCompletion(context.Background(), openai.ChatCompletionRequest{ + Model: "o3-mini", + Messages: []openai.ChatCompletionMessage{{Role: "user", Content: "hi"}}, + MaxTokens: 256, + }) + require.Error(t, err, "sending max_tokens to an o-series model fails before the request leaves the process") + assert.Zero(t, api.requests.Load()) + }) + + t.Run("the stream flag is rejected on the completion call", func(t *testing.T) { + _, err := client.CreateChatCompletion(context.Background(), openai.ChatCompletionRequest{ + Model: "gpt-4o-mini", + Messages: []openai.ChatCompletionMessage{{Role: "user", Content: "hi"}}, + Stream: true, + }) + assert.ErrorIs(t, err, openai.ErrChatCompletionStreamNotSupported) + }) + + t.Run("the default client has no timeout", func(t *testing.T) { + httpClient, ok := openai.DefaultConfig("test-key").HTTPClient.(*http.Client) // pragma: allowlist secret + require.True(t, ok) + assert.Zero(t, httpClient.Timeout, "the provider must install its own timeout") + }) +} + +// The provider is shared between goroutines by the agent runtime; the model +// cache and the configuration were unsynchronised. Run with -race. +func TestOpenAI_ConcurrentUse(t *testing.T) { + api := newOpenAIAPI(t, openAIRoutes(openAIChatSuccess("ok"), func(w http.ResponseWriter, r *http.Request) { + _ = json.NewEncoder(w).Encode(map[string]interface{}{ + "object": "list", + "data": []map[string]interface{}{{"id": "gpt-4o", "object": "model"}}, + }) + })) + p := api.provider(t, nil) + + var wg sync.WaitGroup + errs := make(chan error, 64) + + for i := 0; i < 8; i++ { + wg.Add(1) + go func(i int) { + defer wg.Done() + if _, err := p.Complete(context.Background(), openAITestRequest()); err != nil { + errs <- err + } + if _, err := p.GetModels(context.Background()); err != nil { + errs <- err + } + _ = p.GetConfig() + _ = p.GetStreamingConfig() + // Rebuilds the client while other goroutines are mid-request. + if err := p.SetConfig(map[string]interface{}{ + "model": fmt.Sprintf("gpt-4o-%d", i), + "timeout": time.Duration(5+i) * time.Second, + }); err != nil { + errs <- err + } + }(i) + } + + wg.Wait() + close(errs) + for err := range errs { + t.Errorf("concurrent use failed: %v", err) + } +} diff --git a/pkg/llm/provider.go b/pkg/llm/provider.go index 4c2db46..27c3eed 100644 --- a/pkg/llm/provider.go +++ b/pkg/llm/provider.go @@ -67,7 +67,7 @@ type CompletionRequest struct { // EarlyExit, when set on a streaming completion, is checked after each // chunk. Returning true cancels the remainder of the token stream // (saves SLM decode latency once a complete JSON/tool-call is formed). - EarlyExit EarlyExitFunc `json:"-"` + EarlyExit EarlyExitFunc `json:"-" yaml:"-"` } // CompletionResponse represents a response from completion diff --git a/pkg/llm/resilience.go b/pkg/llm/resilience.go new file mode 100644 index 0000000..e98d59a --- /dev/null +++ b/pkg/llm/resilience.go @@ -0,0 +1,197 @@ +// Copyright (c) 2024 GoLangGraph Team +// +// Licensed under the MIT License. See LICENSE file in the project root for full license information. +// +// Package: GoLangGraph - A powerful Go framework for building AI agent workflows + +package llm + +import ( + "context" + "errors" + "fmt" + "io" + "net/http" + "strconv" + "time" +) + +// Sentinel errors classifying provider failures, so callers can branch with +// errors.Is instead of matching on message text. +var ( + // ErrProviderUnavailable indicates a transient failure: a network error or + // a 5xx response. Retrying may succeed. + ErrProviderUnavailable = errors.New("provider unavailable") + // ErrRateLimited indicates the provider asked the caller to slow down. + ErrRateLimited = errors.New("provider rate limited") + // ErrProviderRequest indicates a permanent client error (4xx). Retrying + // the same request will not help. + ErrProviderRequest = errors.New("invalid provider request") + // ErrProviderAuth indicates rejected credentials. + ErrProviderAuth = errors.New("provider authentication failed") + // ErrResponseTooLarge indicates the provider returned more data than the + // configured limit allows. + ErrResponseTooLarge = errors.New("provider response too large") +) + +// MaxResponseBytes bounds how much of a provider response is buffered, so a +// broken or hostile endpoint cannot exhaust memory. +const MaxResponseBytes int64 = 32 << 20 // 32 MiB + +// ProviderError carries the provider's HTTP status and body alongside a +// classification of whether the request is worth retrying. +type ProviderError struct { + Provider string + StatusCode int + Body string + // RetryAfter is set when the provider supplied a Retry-After header. + RetryAfter time.Duration + kind error +} + +func (e *ProviderError) Error() string { + if e.StatusCode > 0 { + return fmt.Sprintf("%s: %s API error: status %d: %s", e.kind, e.Provider, e.StatusCode, e.Body) + } + return fmt.Sprintf("%s: %s: %s", e.kind, e.Provider, e.Body) +} + +// Unwrap exposes the classification so errors.Is matches the sentinels. +func (e *ProviderError) Unwrap() error { return e.kind } + +// Retryable reports whether another attempt could succeed. +func (e *ProviderError) Retryable() bool { + return errors.Is(e.kind, ErrProviderUnavailable) || errors.Is(e.kind, ErrRateLimited) +} + +// classifyStatus maps an HTTP status onto a sentinel. +func classifyStatus(status int) error { + switch { + case status == http.StatusTooManyRequests: + return ErrRateLimited + case status == http.StatusUnauthorized, status == http.StatusForbidden: + return ErrProviderAuth + case status >= 500: + return ErrProviderUnavailable + case status >= 400: + return ErrProviderRequest + } + return nil +} + +// NewProviderError builds a classified error from an HTTP response. The body is +// read up to a bounded size so an oversized error page cannot exhaust memory. +func NewProviderError(provider string, resp *http.Response) *ProviderError { + body, _ := io.ReadAll(io.LimitReader(resp.Body, 8<<10)) + pe := &ProviderError{ + Provider: provider, + StatusCode: resp.StatusCode, + Body: string(body), + kind: classifyStatus(resp.StatusCode), + } + if pe.kind == nil { + pe.kind = ErrProviderUnavailable + } + if after := resp.Header.Get("Retry-After"); after != "" { + if seconds, err := strconv.Atoi(after); err == nil && seconds >= 0 { + pe.RetryAfter = time.Duration(seconds) * time.Second + } + } + return pe +} + +// NewTransportError wraps a network-level failure as a retryable provider error. +func NewTransportError(provider string, err error) *ProviderError { + return &ProviderError{ + Provider: provider, + Body: err.Error(), + kind: ErrProviderUnavailable, + } +} + +// IsRetryable reports whether an error is worth another attempt. Context +// cancellation is never retryable. +func IsRetryable(err error) bool { + if err == nil { + return false + } + if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) { + return false + } + var pe *ProviderError + if errors.As(err, &pe) { + return pe.Retryable() + } + return false +} + +// retryAfter returns a provider-requested delay, if any. +func retryAfter(err error) time.Duration { + var pe *ProviderError + if errors.As(err, &pe) { + return pe.RetryAfter + } + return 0 +} + +// WithRetry runs fn, retrying transient failures with exponential backoff. +// +// RetryCount and RetryDelay were previously accepted as configuration and then +// ignored, so a transient provider blip failed the whole call. Attempts stop +// early on a permanent error or a canceled context, and a provider-supplied +// Retry-After is honored. +func WithRetry(ctx context.Context, config *ProviderConfig, fn func(ctx context.Context) error) error { + attempts := 0 + delay := time.Second + if config != nil { + attempts = config.RetryCount + if config.RetryDelay > 0 { + delay = config.RetryDelay + } + } + if attempts < 0 { + attempts = 0 + } + + var err error + for attempt := 0; ; attempt++ { + if ctxErr := ctx.Err(); ctxErr != nil { + if err != nil { + return err + } + return ctxErr + } + + err = fn(ctx) + if err == nil { + return nil + } + if attempt >= attempts || !IsRetryable(err) { + return err + } + + wait := delay + if requested := retryAfter(err); requested > wait { + wait = requested + } + + select { + case <-ctx.Done(): + return err + case <-time.After(wait): + } + + // Exponential backoff, capped so a long retry budget cannot stall a + // request for an unbounded time. + delay *= 2 + if delay > 30*time.Second { + delay = 30 * time.Second + } + } +} + +// limitedBody wraps a response body with a size cap, so a provider that streams +// endlessly cannot exhaust memory. +func limitedBody(body io.ReadCloser) io.Reader { + return io.LimitReader(body, MaxResponseBytes) +} diff --git a/pkg/llm/resilience_test.go b/pkg/llm/resilience_test.go new file mode 100644 index 0000000..f9dbb2d --- /dev/null +++ b/pkg/llm/resilience_test.go @@ -0,0 +1,370 @@ +// Copyright (c) 2024 GoLangGraph Team +// +// Licensed under the MIT License. See LICENSE file in the project root for full license information. + +package llm + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "net/http" + "net/http/httptest" + "strings" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// ollamaChatResponse writes one terminal chat frame. +func ollamaChatResponse(w http.ResponseWriter, content string) { + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(map[string]interface{}{ + "model": "test-model", + "created_at": time.Now(), + "message": map[string]string{"role": "assistant", "content": content}, + "done": true, + }) +} + +func testOllama(t *testing.T, endpoint string, mutate func(*ProviderConfig)) *OllamaProvider { + t.Helper() + cfg := DefaultProviderConfig() + cfg.Name = "ollama" + cfg.Type = "ollama" + cfg.Endpoint = endpoint + cfg.Model = "test-model" + cfg.RetryCount = 0 + cfg.RetryDelay = time.Millisecond + cfg.Timeout = 5 * time.Second + if mutate != nil { + mutate(cfg) + } + p, err := NewOllamaProvider(cfg) + require.NoError(t, err) + return p +} + +func simpleRequest() CompletionRequest { + return CompletionRequest{ + Model: "test-model", + Messages: []Message{{Role: "user", Content: "hello"}}, + } +} + +// --------------------------------------------------------------------------- +// Error classification +// --------------------------------------------------------------------------- + +func TestProvider_ClassifiesHTTPStatuses(t *testing.T) { + cases := []struct { + status int + want error + retry bool + }{ + {http.StatusInternalServerError, ErrProviderUnavailable, true}, + {http.StatusBadGateway, ErrProviderUnavailable, true}, + {http.StatusServiceUnavailable, ErrProviderUnavailable, true}, + {http.StatusTooManyRequests, ErrRateLimited, true}, + {http.StatusUnauthorized, ErrProviderAuth, false}, + {http.StatusForbidden, ErrProviderAuth, false}, + {http.StatusBadRequest, ErrProviderRequest, false}, + {http.StatusNotFound, ErrProviderRequest, false}, + } + + for _, tc := range cases { + t.Run(fmt.Sprint(tc.status), func(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + http.Error(w, "upstream said no", tc.status) + })) + defer srv.Close() + + p := testOllama(t, srv.URL, nil) + _, err := p.Complete(context.Background(), simpleRequest()) + require.Error(t, err) + assert.True(t, errors.Is(err, tc.want), "status %d should classify as %v, got %v", tc.status, tc.want, err) + assert.Equal(t, tc.retry, IsRetryable(err), "retryability for status %d", tc.status) + }) + } +} + +// --------------------------------------------------------------------------- +// Retry behavior +// --------------------------------------------------------------------------- + +func TestProvider_RetriesTransientFailures(t *testing.T) { + var calls atomic.Int32 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if calls.Add(1) < 3 { + http.Error(w, "temporarily down", http.StatusServiceUnavailable) + return + } + ollamaChatResponse(w, "recovered") + })) + defer srv.Close() + + p := testOllama(t, srv.URL, func(c *ProviderConfig) { c.RetryCount = 3 }) + resp, err := p.Complete(context.Background(), simpleRequest()) + require.NoError(t, err, "a transient outage must be retried") + assert.Equal(t, "recovered", resp.Choices[0].Message.Content) + assert.EqualValues(t, 3, calls.Load()) +} + +func TestProvider_DoesNotRetryPermanentErrors(t *testing.T) { + var calls atomic.Int32 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + calls.Add(1) + http.Error(w, "bad model", http.StatusBadRequest) + })) + defer srv.Close() + + p := testOllama(t, srv.URL, func(c *ProviderConfig) { c.RetryCount = 5 }) + _, err := p.Complete(context.Background(), simpleRequest()) + require.Error(t, err) + assert.EqualValues(t, 1, calls.Load(), "a 4xx must not be retried") +} + +func TestProvider_ExhaustsRetryBudget(t *testing.T) { + var calls atomic.Int32 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + calls.Add(1) + http.Error(w, "still down", http.StatusInternalServerError) + })) + defer srv.Close() + + p := testOllama(t, srv.URL, func(c *ProviderConfig) { c.RetryCount = 2 }) + _, err := p.Complete(context.Background(), simpleRequest()) + require.Error(t, err) + assert.True(t, errors.Is(err, ErrProviderUnavailable)) + assert.EqualValues(t, 3, calls.Load(), "initial attempt plus RetryCount retries") +} + +func TestProvider_HonorsRetryAfter(t *testing.T) { + var calls atomic.Int32 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if calls.Add(1) == 1 { + w.Header().Set("Retry-After", "1") + http.Error(w, "slow down", http.StatusTooManyRequests) + return + } + ollamaChatResponse(w, "ok") + })) + defer srv.Close() + + p := testOllama(t, srv.URL, func(c *ProviderConfig) { + c.RetryCount = 2 + c.RetryDelay = time.Millisecond + }) + + start := time.Now() + _, err := p.Complete(context.Background(), simpleRequest()) + require.NoError(t, err) + assert.GreaterOrEqual(t, time.Since(start), time.Second, + "a provider-supplied Retry-After must be respected over the configured delay") +} + +// A canceled context must abandon retries immediately. +func TestProvider_CancellationStopsRetries(t *testing.T) { + var calls atomic.Int32 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + calls.Add(1) + http.Error(w, "down", http.StatusInternalServerError) + })) + defer srv.Close() + + p := testOllama(t, srv.URL, func(c *ProviderConfig) { + c.RetryCount = 100 + c.RetryDelay = 50 * time.Millisecond + }) + + ctx, cancel := context.WithTimeout(context.Background(), 120*time.Millisecond) + defer cancel() + + done := make(chan struct{}) + go func() { + defer close(done) + _, _ = p.Complete(ctx, simpleRequest()) + }() + + select { + case <-done: + case <-time.After(10 * time.Second): + t.Fatal("retries did not stop when the context was canceled") + } + assert.Less(t, calls.Load(), int32(20), "retries must stop promptly on cancellation") +} + +// --------------------------------------------------------------------------- +// Network and payload failures +// --------------------------------------------------------------------------- + +func TestProvider_NetworkFailureIsRetryable(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})) + endpoint := srv.URL + srv.Close() // nothing is listening now + + p := testOllama(t, endpoint, nil) + _, err := p.Complete(context.Background(), simpleRequest()) + require.Error(t, err) + assert.True(t, errors.Is(err, ErrProviderUnavailable), "a refused connection is transient, got %v", err) + assert.True(t, IsRetryable(err)) +} + +func TestProvider_MalformedResponseIsReported(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte("{not json at all")) + })) + defer srv.Close() + + p := testOllama(t, srv.URL, nil) + _, err := p.Complete(context.Background(), simpleRequest()) + require.Error(t, err) + assert.Contains(t, err.Error(), "decode") +} + +func TestProvider_TruncatedResponseIsReported(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + // A frame that never completes, then the connection ends. + _, _ = w.Write([]byte(`{"model":"test-model","message":{"role":"assistant","content":"par`)) + })) + defer srv.Close() + + p := testOllama(t, srv.URL, nil) + _, err := p.Complete(context.Background(), simpleRequest()) + require.Error(t, err, "a truncated response must not be reported as success") +} + +func TestProvider_APILevelErrorIsPermanent(t *testing.T) { + var calls atomic.Int32 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + calls.Add(1) + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(map[string]interface{}{"error": "model not found"}) + })) + defer srv.Close() + + p := testOllama(t, srv.URL, func(c *ProviderConfig) { c.RetryCount = 3 }) + _, err := p.Complete(context.Background(), simpleRequest()) + require.Error(t, err) + assert.Contains(t, err.Error(), "model not found") + assert.EqualValues(t, 1, calls.Load(), "an API-level error must not be retried") +} + +// A response body larger than the cap must not be buffered without bound. +func TestProvider_OversizedResponseIsBounded(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + flusher, _ := w.(http.Flusher) + chunk := strings.Repeat("x", 64*1024) + // Stream frames that never set done, far beyond any sane response. + for i := 0; i < 200; i++ { + _, _ = fmt.Fprintf(w, `{"model":"m","message":{"role":"assistant","content":%q},"done":false}`, chunk) + if flusher != nil { + flusher.Flush() + } + } + })) + defer srv.Close() + + p := testOllama(t, srv.URL, nil) + done := make(chan struct{}) + go func() { + defer close(done) + _, _ = p.Complete(context.Background(), simpleRequest()) + }() + + select { + case <-done: + case <-time.After(30 * time.Second): + t.Fatal("an endless response was not bounded") + } +} + +func TestProvider_TimeoutIsReported(t *testing.T) { + release := make(chan struct{}) + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + <-release + })) + defer func() { + close(release) + srv.Close() + }() + + p := testOllama(t, srv.URL, func(c *ProviderConfig) { c.Timeout = 100 * time.Millisecond }) + _, err := p.Complete(context.Background(), simpleRequest()) + require.Error(t, err, "a hung provider must not block indefinitely") +} + +// --------------------------------------------------------------------------- +// Manager behavior +// --------------------------------------------------------------------------- + +func TestProviderManager_ConcurrentUse(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + ollamaChatResponse(w, "ok") + })) + defer srv.Close() + + mgr := NewProviderManager() + require.NoError(t, mgr.RegisterProvider("ollama", testOllama(t, srv.URL, nil))) + + var wg sync.WaitGroup + errs := make([]error, 16) + for i := 0; i < 16; i++ { + wg.Add(1) + go func(i int) { + defer wg.Done() + _, errs[i] = mgr.Complete(context.Background(), "ollama", simpleRequest()) + _ = mgr.ListProviders() + _ = mgr.HealthCheck(context.Background()) + }(i) + } + wg.Wait() + require.NoError(t, errors.Join(errs...)) +} + +func TestProviderManager_UnknownProvider(t *testing.T) { + mgr := NewProviderManager() + _, err := mgr.Complete(context.Background(), "nope", simpleRequest()) + require.Error(t, err) +} + +func TestProviderManager_RejectsDuplicateRegistration(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + ollamaChatResponse(w, "ok") + })) + defer srv.Close() + + mgr := NewProviderManager() + require.NoError(t, mgr.RegisterProvider("p", testOllama(t, srv.URL, nil))) + assert.Error(t, mgr.RegisterProvider("p", testOllama(t, srv.URL, nil))) +} + +// HealthCheck must report per-provider status rather than failing as a whole. +func TestProviderManager_HealthCheckReportsPerProvider(t *testing.T) { + good := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _ = json.NewEncoder(w).Encode(map[string]interface{}{"models": []interface{}{}}) + })) + defer good.Close() + + bad := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + http.Error(w, "down", http.StatusInternalServerError) + })) + badURL := bad.URL + bad.Close() + + mgr := NewProviderManager() + require.NoError(t, mgr.RegisterProvider("good", testOllama(t, good.URL, nil))) + require.NoError(t, mgr.RegisterProvider("bad", testOllama(t, badURL, nil))) + + health := mgr.HealthCheck(context.Background()) + assert.NoError(t, health["good"]) + assert.Error(t, health["bad"], "an unreachable provider must be reported unhealthy") +} diff --git a/pkg/llm/streaming_test.go b/pkg/llm/streaming_test.go index bdac1d2..64d39db 100644 --- a/pkg/llm/streaming_test.go +++ b/pkg/llm/streaming_test.go @@ -8,6 +8,9 @@ package llm import ( "context" + "encoding/json" + "fmt" + "net/http" "strings" "testing" "time" @@ -16,6 +19,34 @@ import ( "github.com/stretchr/testify/require" ) +// geminiStreamingAPI is a fake Generative Language API that streams several +// chunks, so streaming tests exercise the real SSE parsing path. +func geminiStreamingAPI(t *testing.T, parts ...string) *geminiAPI { + t.Helper() + if len(parts) == 0 { + parts = []string{"one ", "two ", "three"} + } + return newGeminiAPI(t, func(w http.ResponseWriter, r *http.Request) { + if strings.Contains(r.URL.Path, "streamGenerateContent") { + w.Header().Set("Content-Type", "text/event-stream") + flusher, _ := w.(http.Flusher) + for _, part := range parts { + frame, _ := json.Marshal(map[string]interface{}{ + "candidates": []map[string]interface{}{{ + "content": map[string]interface{}{"parts": []map[string]string{{"text": part}}}, + }}, + }) + _, _ = fmt.Fprintf(w, "data: %s\n\n", frame) + if flusher != nil { + flusher.Flush() + } + } + return + } + geminiSuccess(strings.Join(parts, ""))(w, r) + }) +} + func TestStreamingConfig(t *testing.T) { config := DefaultStreamingConfig() assert.False(t, config.Enabled) @@ -167,13 +198,23 @@ func TestProviderManagerStreaming(t *testing.T) { } func TestStreamingModesWithGemini(t *testing.T) { - config := DefaultProviderConfig() - config.Type = "gemini" - config.APIKey = "test-key" // pragma: allowlist secret - config.Model = "gemini-pro" - - provider, err := NewGeminiProvider(config) - require.NoError(t, err) + // Backed by a fake Generative Language API: the provider builds real + // requests and parses real responses. Before, this passed against a + // hardcoded reply that never left the process. + api := newGeminiAPI(t, func(w http.ResponseWriter, r *http.Request) { + if strings.Contains(r.URL.Path, "streamGenerateContent") { + w.Header().Set("Content-Type", "text/event-stream") + frame, _ := json.Marshal(map[string]interface{}{ + "candidates": []map[string]interface{}{{ + "content": map[string]interface{}{"parts": []map[string]string{{"text": "streamed"}}}, + }}, + }) + _, _ = fmt.Fprintf(w, "data: %s\n\n", frame) + return + } + geminiSuccess("complete")(w, r) + }) + provider := api.provider(t, func(c *ProviderConfig) { c.Model = "gemini-pro" }) ctx := context.Background() req := CompletionRequest{ @@ -209,11 +250,10 @@ func TestStreamingModesWithGemini(t *testing.T) { func TestStreamingCallbackWithGemini(t *testing.T) { config := DefaultProviderConfig() config.Type = "gemini" - config.APIKey = "test-key" // pragma: allowlist secret - config.Model = "gemini-pro" + _ = config - provider, err := NewGeminiProvider(config) - require.NoError(t, err) + api := geminiStreamingAPI(t, "1 ", "2 ", "3 ", "4 ", "5") + provider := api.provider(t, func(c *ProviderConfig) { c.Model = "gemini-pro" }) ctx := context.Background() req := CompletionRequest{ @@ -231,7 +271,7 @@ func TestStreamingCallbackWithGemini(t *testing.T) { } // Test streaming with callback - err = provider.CompleteStreamWithMode(ctx, req, callback, StreamModeForced) + err := provider.CompleteStreamWithMode(ctx, req, callback, StreamModeForced) assert.NoError(t, err) assert.Greater(t, len(chunks), 1) // Should receive multiple chunks @@ -326,16 +366,12 @@ func TestRealisticStreamingScenario(t *testing.T) { // Create provider manager pm := NewProviderManager() - // Register Gemini provider (using mock) - config := DefaultProviderConfig() - config.Type = "gemini" - config.APIKey = "test-key" // pragma: allowlist secret - config.Model = "gemini-pro" - - provider, err := NewGeminiProvider(config) - require.NoError(t, err) + // Register a Gemini provider backed by a fake API, so the streaming path + // under test is the real one. + api := geminiStreamingAPI(t, "Once ", "upon ", "a ", "time") + provider := api.provider(t, func(c *ProviderConfig) { c.Model = "gemini-pro" }) - err = pm.RegisterProvider("gemini", provider) + err := pm.RegisterProvider("gemini", provider) require.NoError(t, err) // Enable streaming diff --git a/pkg/persistence/checkpointer.go b/pkg/persistence/checkpointer.go index 365d25f..7388a8b 100644 --- a/pkg/persistence/checkpointer.go +++ b/pkg/persistence/checkpointer.go @@ -10,6 +10,11 @@ import ( "context" "encoding/json" "fmt" + "os" + "path/filepath" + "regexp" + "sort" + "strings" "sync" "time" @@ -241,8 +246,15 @@ func (c *FileCheckpointer) Save(ctx context.Context, checkpoint *Checkpoint) err c.mu.Lock() defer c.mu.Unlock() + if err := validateID("thread ID", checkpoint.ThreadID); err != nil { + return err + } + if err := validateID("checkpoint ID", checkpoint.ID); err != nil { + return err + } + // Create directory structure - threadDir := fmt.Sprintf("%s/%s", c.basePath, checkpoint.ThreadID) + threadDir := filepath.Join(c.basePath, checkpoint.ThreadID) if err := ensureDir(threadDir); err != nil { return fmt.Errorf("failed to create thread directory: %w", err) } @@ -254,7 +266,7 @@ func (c *FileCheckpointer) Save(ctx context.Context, checkpoint *Checkpoint) err } // Write to file - filePath := fmt.Sprintf("%s/%s.json", threadDir, checkpoint.ID) + filePath := filepath.Join(threadDir, checkpoint.ID+".json") if err := writeFile(filePath, data); err != nil { return fmt.Errorf("failed to write checkpoint file: %w", err) } @@ -267,7 +279,14 @@ func (c *FileCheckpointer) Load(ctx context.Context, threadID, checkpointID stri c.mu.RLock() defer c.mu.RUnlock() - filePath := fmt.Sprintf("%s/%s/%s.json", c.basePath, threadID, checkpointID) + if err := validateID("thread ID", threadID); err != nil { + return nil, err + } + if err := validateID("checkpoint ID", checkpointID); err != nil { + return nil, err + } + + filePath := filepath.Join(c.basePath, threadID, checkpointID+".json") data, err := readFile(filePath) if err != nil { @@ -287,7 +306,11 @@ func (c *FileCheckpointer) List(ctx context.Context, threadID string) ([]*Checkp c.mu.RLock() defer c.mu.RUnlock() - threadDir := fmt.Sprintf("%s/%s", c.basePath, threadID) + if err := validateID("thread ID", threadID); err != nil { + return nil, err + } + + threadDir := filepath.Join(c.basePath, threadID) files, err := listFiles(threadDir, ".json") if err != nil { @@ -296,7 +319,7 @@ func (c *FileCheckpointer) List(ctx context.Context, threadID string) ([]*Checkp var metadata []*CheckpointMetadata for _, file := range files { - filePath := fmt.Sprintf("%s/%s", threadDir, file) + filePath := filepath.Join(threadDir, file) data, err := readFile(filePath) if err != nil { @@ -328,7 +351,14 @@ func (c *FileCheckpointer) Delete(ctx context.Context, threadID, checkpointID st c.mu.Lock() defer c.mu.Unlock() - filePath := fmt.Sprintf("%s/%s/%s.json", c.basePath, threadID, checkpointID) + if err := validateID("thread ID", threadID); err != nil { + return err + } + if err := validateID("checkpoint ID", checkpointID); err != nil { + return err + } + + filePath := filepath.Join(c.basePath, threadID, checkpointID+".json") if err := deleteFile(filePath); err != nil { return fmt.Errorf("failed to delete checkpoint file: %w", err) @@ -482,28 +512,83 @@ func (tt *TimeTravel) FindCheckpointByNode(ctx context.Context, threadID, nodeID return latest, nil } -// Placeholder functions for file operations (would be implemented with actual file I/O) -func ensureDir(path string) error { - // Implementation would create directory if it doesn't exist +// File operations backing FileCheckpointer. +// safeIDPattern bounds the characters allowed in identifiers that become path +// components. Thread and checkpoint IDs arrive from API clients, so without +// this a value such as "../../etc" would escape the checkpoint directory. +var safeIDPattern = regexp.MustCompile(`^[A-Za-z0-9._-]{1,190}$`) + +// validateID rejects identifiers that are unsafe to use as a path component. +func validateID(kind, id string) error { + if id == "" { + return fmt.Errorf("%s must not be empty", kind) + } + if strings.Contains(id, "/") || strings.Contains(id, `\\`) || strings.Contains(id, "..") { + return fmt.Errorf("%s %q contains path separators", kind, id) + } + if !safeIDPattern.MatchString(id) { + return fmt.Errorf("%s %q must match %s", kind, id, safeIDPattern.String()) + } return nil } +func ensureDir(path string) error { + return os.MkdirAll(path, 0o750) +} + func writeFile(path string, data []byte) error { - // Implementation would write data to file - return nil + // Write to a temporary file and rename so a crash mid-write cannot leave a + // truncated checkpoint behind. + dir := filepath.Dir(path) + tmp, err := os.CreateTemp(dir, ".checkpoint-*.tmp") + if err != nil { + return err + } + tmpName := tmp.Name() + defer func() { + _ = os.Remove(tmpName) + }() + + if _, err := tmp.Write(data); err != nil { + _ = tmp.Close() + return err + } + if err := tmp.Sync(); err != nil { + _ = tmp.Close() + return err + } + if err := tmp.Close(); err != nil { + return err + } + if err := os.Chmod(tmpName, 0o640); err != nil { + return err + } + return os.Rename(tmpName, path) } func readFile(path string) ([]byte, error) { - // Implementation would read data from file - return []byte{}, nil + return os.ReadFile(path) // #nosec G304 -- path is composed from validated checkpoint identifiers } func listFiles(dir, extension string) ([]string, error) { - // Implementation would list files with given extension in directory - return []string{}, nil + entries, err := os.ReadDir(dir) + if err != nil { + return nil, err + } + var files []string + for _, entry := range entries { + if entry.IsDir() { + continue + } + if extension != "" && filepath.Ext(entry.Name()) != extension { + continue + } + files = append(files, entry.Name()) + } + sort.Strings(files) + return files, nil } func deleteFile(path string) error { - // Implementation would delete file - return nil + return os.Remove(path) } diff --git a/pkg/persistence/database.go b/pkg/persistence/database.go index 7fb506d..e094b22 100644 --- a/pkg/persistence/database.go +++ b/pkg/persistence/database.go @@ -11,7 +11,9 @@ import ( "database/sql" "encoding/json" "fmt" + "strconv" "strings" + "sync" "time" "github.com/go-redis/redis/v8" @@ -49,6 +51,11 @@ type DatabaseConfig struct { MaxIdleConns int `json:"max_idle_conns"` MaxLifetime string `json:"max_lifetime"` + // CheckpointTTL bounds how long a checkpoint survives in stores that expire + // keys (Redis). Empty means the 24h default. Set it to a duration string + // such as "168h"; "0" disables expiry entirely. + CheckpointTTL string `json:"checkpoint_ttl"` + // Vector-specific configuration VectorDimension int `json:"vector_dimension"` VectorMetric string `json:"vector_metric"` // "cosine", "euclidean", "dot_product" @@ -149,15 +156,28 @@ func (p *PostgresConnection) Connect() error { } if p.config.MaxLifetime != "" { - if duration, err := time.ParseDuration(p.config.MaxLifetime); err == nil { - db.SetConnMaxLifetime(duration) + // A typo'd duration used to be swallowed silently, leaving connections + // with no lifetime cap at all -- the opposite of what the operator + // configured. Fail loudly instead. + duration, perr := time.ParseDuration(p.config.MaxLifetime) + if perr != nil { + _ = db.Close() + return fmt.Errorf("invalid max_lifetime %q: %w", p.config.MaxLifetime, perr) } + db.SetConnMaxLifetime(duration) } else { db.SetConnMaxLifetime(5 * time.Minute) // Default } p.db = db - return p.Ping() + if err := p.Ping(); err != nil { + // Ping failing leaves an open *sql.DB (and its pool goroutines) behind + // unless we close it; callers only see the error and drop the object. + _ = db.Close() + p.db = nil + return err + } + return nil } // Ping tests the database connection @@ -169,6 +189,9 @@ func (p *PostgresConnection) Ping() error { ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() + if p.db == nil { + return fmt.Errorf("database connection is not open") + } return p.db.PingContext(ctx) } @@ -192,20 +215,177 @@ func (p *PostgresConnection) GetConfig() *DatabaseConfig { // ExecuteQuery executes a query without returning results func (p *PostgresConnection) ExecuteQuery(ctx context.Context, query string, args ...interface{}) error { + if p.db == nil { + return fmt.Errorf("database connection is not open") + } _, err := p.db.ExecContext(ctx, query, args...) return err } // QueryRow executes a query that returns a single row func (p *PostgresConnection) QueryRow(ctx context.Context, query string, args ...interface{}) interface{} { + if p.db == nil { + return nil + } return p.db.QueryRowContext(ctx, query, args...) } // QueryRows executes a query that returns multiple rows func (p *PostgresConnection) QueryRows(ctx context.Context, query string, args ...interface{}) (interface{}, error) { + if p.db == nil { + return nil, fmt.Errorf("database connection is not open") + } return p.db.QueryContext(ctx, query, args...) } +// Exec runs a statement and returns its sql.Result. +// +// ExecuteQuery throws the result away, which made it impossible for callers to +// tell "deleted one row" from "matched nothing" -- see Delete below. +func (p *PostgresConnection) Exec(ctx context.Context, query string, args ...interface{}) (sql.Result, error) { + if p.db == nil { + return nil, fmt.Errorf("database connection is not open") + } + return p.db.ExecContext(ctx, query, args...) +} + +// WithTx runs fn inside a transaction, committing on success and rolling back +// on any error or panic. +func (p *PostgresConnection) WithTx(ctx context.Context, fn func(*sql.Tx) error) (err error) { + if p.db == nil { + return fmt.Errorf("database connection is not open") + } + + tx, err := p.db.BeginTx(ctx, nil) + if err != nil { + return fmt.Errorf("failed to begin transaction: %w", err) + } + + defer func() { + if r := recover(); r != nil { + _ = tx.Rollback() + panic(r) + } + if err != nil { + // Rollback error is deliberately not surfaced: the caller needs the + // original failure, and a rollback after a failed statement is + // frequently a no-op the driver reports as ErrTxDone. + _ = tx.Rollback() + } + }() + + if err = fn(tx); err != nil { + return err + } + + if err = tx.Commit(); err != nil { + return fmt.Errorf("failed to commit transaction: %w", err) + } + return nil +} + +// asSQLRow converts a DatabaseConnection result to *sql.Row. +// +// DatabaseConnection is a public interface returning interface{}, so an +// implementation other than PostgresConnection previously caused a panic at +// the unchecked type assertion. Returning an error keeps a custom or absent +// backend from crashing the process. +func asSQLRow(v interface{}) (*sql.Row, error) { + row, ok := v.(*sql.Row) + if !ok || row == nil { + return nil, fmt.Errorf("database connection returned %T, want *sql.Row", v) + } + return row, nil +} + +// asSQLRows is the multi-row counterpart of asSQLRow. +// +// The row-iterating code used to write `rows.(*sql.Rows)` inline -- a +// single-value type assertion that panics rather than erroring. QueryRows is +// declared on the public DatabaseConnection interface as returning interface{}, +// so any implementation other than PostgresConnection crashed the process. +func asSQLRows(v interface{}) (*sql.Rows, error) { + rows, ok := v.(*sql.Rows) + if !ok || rows == nil { + return nil, fmt.Errorf("database connection returned %T, want *sql.Rows", v) + } + return rows, nil +} + +// decodeJSONMap unmarshals a JSONB column into a map, tolerating SQL NULL and +// the JSON literal null. +// +// Every caller previously ran json.Unmarshal on the raw bytes, so a row with a +// NULL metadata column -- which the schema permits, and which any row written +// by another tool or an older release may well have -- failed with "unexpected +// end of JSON input". That broke Load *and* List, and a broken List breaks +// Latest(), i.e. resuming a thread at all. +func decodeJSONMap(data []byte, target *map[string]interface{}) error { + if len(data) == 0 || string(data) == "null" { + *target = map[string]interface{}{} + return nil + } + if err := json.Unmarshal(data, target); err != nil { + return err + } + if *target == nil { + *target = map[string]interface{}{} + } + return nil +} + +// encodeVector renders a float slice as a pgvector literal ("[1,2,3]"). +// +// The RAG methods used to hand []float64 straight to database/sql, which +// rejects it with "unsupported type []float64, a slice of float64" -- so +// SaveDocument and the vector branch of SearchDocuments could never succeed. +func encodeVector(v []float64) string { + var b strings.Builder + b.WriteByte('[') + for i, f := range v { + if i > 0 { + b.WriteByte(',') + } + b.WriteString(strconv.FormatFloat(f, 'f', -1, 64)) + } + b.WriteByte(']') + return b.String() +} + +// decodeVector parses a pgvector literal back into a float slice. Returns nil +// for SQL NULL so an absent embedding stays absent rather than becoming []. +func decodeVector(raw interface{}) ([]float64, error) { + var s string + switch v := raw.(type) { + case nil: + return nil, nil + case []byte: + s = string(v) + case string: + s = v + default: + return nil, fmt.Errorf("unexpected embedding column type %T", raw) + } + + s = strings.TrimSpace(s) + s = strings.TrimPrefix(s, "[") + s = strings.TrimSuffix(s, "]") + if s == "" { + return nil, nil + } + + parts := strings.Split(s, ",") + out := make([]float64, 0, len(parts)) + for _, p := range parts { + f, err := strconv.ParseFloat(strings.TrimSpace(p), 64) + if err != nil { + return nil, fmt.Errorf("failed to parse embedding component %q: %w", p, err) + } + out = append(out, f) + } + return out, nil +} + // PostgresCheckpointer implements database-based checkpointing with PostgreSQL type PostgresCheckpointer struct { conn *PostgresConnection @@ -228,6 +408,9 @@ func NewPostgresCheckpointer(config *DatabaseConfig) (*PostgresCheckpointer, err // Initialize schema if err := checkpointer.initSchema(); err != nil { + // The connection pool was already open at this point and used to be + // abandoned here, leaking sockets and goroutines on every failed start. + _ = conn.Close() return nil, fmt.Errorf("failed to initialize schema: %w", err) } @@ -397,7 +580,23 @@ func (p *PostgresCheckpointer) Save(ctx context.Context, checkpoint *Checkpoint) return fmt.Errorf("failed to marshal metadata: %w", err) } - query := ` + // checkpoints.thread_id carries a FOREIGN KEY to threads(id), but nothing in + // the Checkpointer interface creates threads -- so every Save against a + // thread that had not been registered out-of-band failed with + // "violates foreign key constraint checkpoints_thread_id_fkey". The + // in-memory and file checkpointers have no such requirement, so the + // PostgreSQL backend was not usable as a drop-in Checkpointer at all. + // + // Registering the parent thread here makes Save self-sufficient. Both + // statements run in one transaction so a failed checkpoint write cannot + // leave an orphan thread row behind. + const ensureThread = ` + INSERT INTO threads (id, created_at, updated_at) + VALUES ($1, NOW(), NOW()) + ON CONFLICT (id) DO NOTHING + ` + + const upsertCheckpoint = ` INSERT INTO checkpoints (id, thread_id, state_data, metadata, created_at, node_id, step_id) VALUES ($1, $2, $3, $4, $5, $6, $7) ON CONFLICT (id) DO UPDATE SET @@ -408,15 +607,23 @@ func (p *PostgresCheckpointer) Save(ctx context.Context, checkpoint *Checkpoint) step_id = EXCLUDED.step_id ` - err = p.conn.ExecuteQuery(ctx, query, - checkpoint.ID, - checkpoint.ThreadID, - stateData, - metadataData, - checkpoint.CreatedAt, - checkpoint.NodeID, - checkpoint.StepID, - ) + err = p.conn.WithTx(ctx, func(tx *sql.Tx) error { + if _, txErr := tx.ExecContext(ctx, ensureThread, checkpoint.ThreadID); txErr != nil { + return fmt.Errorf("failed to register thread: %w", txErr) + } + if _, txErr := tx.ExecContext(ctx, upsertCheckpoint, + checkpoint.ID, + checkpoint.ThreadID, + stateData, + metadataData, + checkpoint.CreatedAt, + checkpoint.NodeID, + checkpoint.StepID, + ); txErr != nil { + return txErr + } + return nil + }) if err != nil { return fmt.Errorf("failed to save checkpoint: %w", err) @@ -438,19 +645,26 @@ func (p *PostgresCheckpointer) Load(ctx context.Context, threadID, checkpointID WHERE thread_id = $1 AND id = $2 ` - row := p.conn.QueryRow(ctx, query, threadID, checkpointID).(*sql.Row) + row, err := asSQLRow(p.conn.QueryRow(ctx, query, threadID, checkpointID)) + if err != nil { + return nil, err + } var checkpoint Checkpoint var stateData, metadataData []byte + // node_id and step_id are nullable in the schema; scanning a NULL straight + // into string/int fails with "converting NULL to string is unsupported". + var nodeID sql.NullString + var stepID sql.NullInt64 - err := row.Scan( + err = row.Scan( &checkpoint.ID, &checkpoint.ThreadID, &stateData, &metadataData, &checkpoint.CreatedAt, - &checkpoint.NodeID, - &checkpoint.StepID, + &nodeID, + &stepID, ) if err != nil { @@ -460,6 +674,9 @@ func (p *PostgresCheckpointer) Load(ctx context.Context, threadID, checkpointID return nil, fmt.Errorf("failed to load checkpoint: %w", err) } + checkpoint.NodeID = nodeID.String + checkpoint.StepID = int(stepID.Int64) + // Unmarshal state var state core.BaseState if err := json.Unmarshal(stateData, &state); err != nil { @@ -468,7 +685,7 @@ func (p *PostgresCheckpointer) Load(ctx context.Context, threadID, checkpointID checkpoint.State = &state // Unmarshal metadata - if err := json.Unmarshal(metadataData, &checkpoint.Metadata); err != nil { + if err := decodeJSONMap(metadataData, &checkpoint.Metadata); err != nil { return nil, fmt.Errorf("failed to unmarshal metadata: %w", err) } @@ -484,38 +701,54 @@ func (p *PostgresCheckpointer) List(ctx context.Context, threadID string) ([]*Ch ORDER BY created_at DESC ` - rows, err := p.conn.QueryRows(ctx, query, threadID) + raw, err := p.conn.QueryRows(ctx, query, threadID) if err != nil { return nil, fmt.Errorf("failed to list checkpoints: %w", err) } - defer rows.(*sql.Rows).Close() + rows, err := asSQLRows(raw) + if err != nil { + return nil, err + } + defer func() { _ = rows.Close() }() - var checkpoints []*CheckpointMetadata - for rows.(*sql.Rows).Next() { + checkpoints := []*CheckpointMetadata{} + for rows.Next() { var checkpoint CheckpointMetadata var metadataData []byte + var nodeID sql.NullString + var stepID sql.NullInt64 - err := rows.(*sql.Rows).Scan( + err := rows.Scan( &checkpoint.ID, &checkpoint.ThreadID, &metadataData, &checkpoint.CreatedAt, - &checkpoint.NodeID, - &checkpoint.StepID, + &nodeID, + &stepID, ) if err != nil { return nil, fmt.Errorf("failed to scan checkpoint: %w", err) } + checkpoint.NodeID = nodeID.String + checkpoint.StepID = int(stepID.Int64) + // Unmarshal metadata - if err := json.Unmarshal(metadataData, &checkpoint.Metadata); err != nil { + if err := decodeJSONMap(metadataData, &checkpoint.Metadata); err != nil { return nil, fmt.Errorf("failed to unmarshal metadata: %w", err) } checkpoints = append(checkpoints, &checkpoint) } + // Without this check a connection that drops mid-iteration returns a + // silently truncated list and a nil error -- the caller cannot tell a + // partial result from a complete one. + if err := rows.Err(); err != nil { + return nil, fmt.Errorf("failed to iterate checkpoints: %w", err) + } + return checkpoints, nil } @@ -523,11 +756,23 @@ func (p *PostgresCheckpointer) List(ctx context.Context, threadID string) ([]*Ch func (p *PostgresCheckpointer) Delete(ctx context.Context, threadID, checkpointID string) error { query := `DELETE FROM checkpoints WHERE thread_id = $1 AND id = $2` - err := p.conn.ExecuteQuery(ctx, query, threadID, checkpointID) + res, err := p.conn.Exec(ctx, query, threadID, checkpointID) if err != nil { return fmt.Errorf("failed to delete checkpoint: %w", err) } + // Deleting a checkpoint that is not there used to report success, so a + // typo'd or already-collected ID looked like a completed deletion. The + // memory and file checkpointers both return an error here; matching them + // keeps the Checkpointer contract the same across backends. + affected, err := res.RowsAffected() + if err != nil { + return fmt.Errorf("failed to confirm checkpoint deletion: %w", err) + } + if affected == 0 { + return fmt.Errorf("checkpoint %s not found in thread %s", checkpointID, threadID) + } + return nil } @@ -544,6 +789,18 @@ func (p *PostgresCheckpointer) SaveDocument(ctx context.Context, doc *Document) return fmt.Errorf("RAG is not enabled") } + if doc == nil { + return fmt.Errorf("cannot save a nil document") + } + + // doc.Metadata used to be handed to database/sql as a bare map, which the + // driver rejects with "unsupported type map[string]interface {}, a map". + // SaveDocument therefore failed 100% of the time; it had never been run. + metadataData, err := json.Marshal(doc.Metadata) + if err != nil { + return fmt.Errorf("failed to marshal document metadata: %w", err) + } + var query string var args []interface{} @@ -557,7 +814,9 @@ func (p *PostgresCheckpointer) SaveDocument(ctx context.Context, doc *Document) embedding = EXCLUDED.embedding, updated_at = EXCLUDED.updated_at ` - args = []interface{}{doc.ID, doc.ThreadID, doc.Content, doc.Metadata, doc.Embedding, doc.CreatedAt, doc.UpdatedAt} + // Likewise []float64 is not a driver value; pgvector accepts its text + // literal form and casts it to the column type. + args = []interface{}{doc.ID, doc.ThreadID, doc.Content, metadataData, encodeVector(doc.Embedding), doc.CreatedAt, doc.UpdatedAt} } else { query = ` INSERT INTO documents (id, thread_id, content, metadata, created_at, updated_at) @@ -567,10 +826,25 @@ func (p *PostgresCheckpointer) SaveDocument(ctx context.Context, doc *Document) metadata = EXCLUDED.metadata, updated_at = EXCLUDED.updated_at ` - args = []interface{}{doc.ID, doc.ThreadID, doc.Content, doc.Metadata, doc.CreatedAt, doc.UpdatedAt} + args = []interface{}{doc.ID, doc.ThreadID, doc.Content, metadataData, doc.CreatedAt, doc.UpdatedAt} } - return p.conn.ExecuteQuery(ctx, query, args...) + // documents.thread_id references threads(id); register the parent for the + // same reason Save does, so a document can be stored for a thread that has + // not been created out-of-band. + return p.conn.WithTx(ctx, func(tx *sql.Tx) error { + if doc.ThreadID != "" { + if _, txErr := tx.ExecContext(ctx, + `INSERT INTO threads (id, created_at, updated_at) VALUES ($1, NOW(), NOW()) ON CONFLICT (id) DO NOTHING`, + doc.ThreadID); txErr != nil { + return fmt.Errorf("failed to register thread: %w", txErr) + } + } + if _, txErr := tx.ExecContext(ctx, query, args...); txErr != nil { + return fmt.Errorf("failed to save document: %w", txErr) + } + return nil + }) } // SearchDocuments performs similarity search on documents @@ -587,10 +861,14 @@ func (p *PostgresCheckpointer) SearchDocuments(ctx context.Context, threadID str SELECT id, thread_id, content, metadata, embedding, created_at, updated_at FROM documents WHERE thread_id = $1 - ORDER BY embedding <-> $2 + ORDER BY embedding <-> $2::vector LIMIT $3 ` - args = []interface{}{threadID, queryEmbedding, limit} + // The raw []float64 the caller passes is not a valid driver value, so + // every vector similarity search failed with "unsupported type + // []float64, a slice of float64". Send the pgvector text literal and + // cast it server-side. + args = []interface{}{threadID, encodeVector(queryEmbedding), limit} } else { // Fallback to text search query = ` @@ -603,39 +881,60 @@ func (p *PostgresCheckpointer) SearchDocuments(ctx context.Context, threadID str args = []interface{}{threadID, limit} } - rows, err := p.conn.QueryRows(ctx, query, args...) + raw, err := p.conn.QueryRows(ctx, query, args...) if err != nil { return nil, fmt.Errorf("failed to search documents: %w", err) } - defer rows.(*sql.Rows).Close() + rows, err := asSQLRows(raw) + if err != nil { + return nil, err + } + defer func() { _ = rows.Close() }() - var documents []*Document - for rows.(*sql.Rows).Next() { + documents := []*Document{} + for rows.Next() { var doc Document var metadataData []byte - var embedding interface{} + // thread_id is nullable on documents. + var docThreadID sql.NullString if p.config.Type == DatabaseTypePgVector { - err := rows.(*sql.Rows).Scan(&doc.ID, &doc.ThreadID, &doc.Content, &metadataData, &embedding, &doc.CreatedAt, &doc.UpdatedAt) + var embedding interface{} + err := rows.Scan(&doc.ID, &docThreadID, &doc.Content, &metadataData, &embedding, &doc.CreatedAt, &doc.UpdatedAt) if err != nil { return nil, fmt.Errorf("failed to scan document: %w", err) } - // Handle embedding conversion if needed + // The stored embedding used to be scanned and then dropped on the + // floor behind a "handle conversion if needed" comment, so every + // document read back had a nil Embedding regardless of what was in + // the column. Decode it properly. + doc.Embedding, err = decodeVector(embedding) + if err != nil { + return nil, fmt.Errorf("failed to decode embedding for document %s: %w", doc.ID, err) + } } else { - err := rows.(*sql.Rows).Scan(&doc.ID, &doc.ThreadID, &doc.Content, &metadataData, &doc.CreatedAt, &doc.UpdatedAt) + err := rows.Scan(&doc.ID, &docThreadID, &doc.Content, &metadataData, &doc.CreatedAt, &doc.UpdatedAt) if err != nil { return nil, fmt.Errorf("failed to scan document: %w", err) } } + doc.ThreadID = docThreadID.String + // Unmarshal metadata - if err := json.Unmarshal(metadataData, &doc.Metadata); err != nil { + if err := decodeJSONMap(metadataData, &doc.Metadata); err != nil { return nil, fmt.Errorf("failed to unmarshal metadata: %w", err) } documents = append(documents, &doc) } + // See List: an unchecked rows.Err() turns a mid-iteration failure into a + // silently short result set. + if err := rows.Err(); err != nil { + return nil, fmt.Errorf("failed to iterate documents: %w", err) + } + return documents, nil } @@ -647,6 +946,45 @@ type RedisCheckpointer struct { ttl time.Duration } +// redisKeySegment escapes an identifier for safe use inside a colon-delimited +// Redis key. +// +// Keys were built with a plain fmt.Sprintf("checkpoint:%s:%s", threadID, id), +// so any identifier containing a colon made distinct checkpoints collide: +// thread "a:b" + checkpoint "c" and thread "a" + checkpoint "b:c" both produced +// "checkpoint:a:b:c". One thread then read and overwrote another thread's +// state. Thread IDs are routinely derived from user or session identifiers, so +// this was a cross-tenant data leak, not just an oddity. +// +// Escaping only ':' and the escape character itself leaves keys byte-identical +// for the ordinary identifiers that contain neither, so existing data stays +// readable. +func redisKeySegment(s string) string { + if !strings.ContainsAny(s, ":%") { + return s + } + var b strings.Builder + for _, r := range s { + switch r { + case ':': + b.WriteString("%3A") + case '%': + b.WriteString("%25") + default: + b.WriteRune(r) + } + } + return b.String() +} + +func redisCheckpointKey(threadID, checkpointID string) string { + return fmt.Sprintf("checkpoint:%s:%s", redisKeySegment(threadID), redisKeySegment(checkpointID)) +} + +func redisThreadIndexKey(threadID string) string { + return fmt.Sprintf("thread:%s:checkpoints", redisKeySegment(threadID)) +} + // NewRedisCheckpointer creates a new Redis checkpointer func NewRedisCheckpointer(config *DatabaseConfig) (*RedisCheckpointer, error) { client := redis.NewClient(&redis.Options{ @@ -660,42 +998,73 @@ func NewRedisCheckpointer(config *DatabaseConfig) (*RedisCheckpointer, error) { defer cancel() if err := client.Ping(ctx).Err(); err != nil { + // The client owns a connection pool and background goroutines; the + // failure path used to drop it without closing, leaking both on every + // unsuccessful connection attempt. + _ = client.Close() return nil, fmt.Errorf("failed to connect to Redis: %w", err) } + // Checkpoints expire after this long. The value was previously hard-coded + // with no way to change it, so every deployment silently lost its + // checkpoints after 24 hours. + ttl := 24 * time.Hour + if config.CheckpointTTL != "" { + parsed, err := time.ParseDuration(config.CheckpointTTL) + if err != nil { + _ = client.Close() + return nil, fmt.Errorf("invalid checkpoint_ttl %q: %w", config.CheckpointTTL, err) + } + ttl = parsed + } + return &RedisCheckpointer{ client: client, config: config, logger: logrus.New(), - ttl: 24 * time.Hour, // Default TTL + ttl: ttl, }, nil } // Save saves a checkpoint to Redis func (r *RedisCheckpointer) Save(ctx context.Context, checkpoint *Checkpoint) error { + if checkpoint == nil { + return fmt.Errorf("cannot save a nil checkpoint") + } + data, err := json.Marshal(checkpoint) if err != nil { return fmt.Errorf("failed to marshal checkpoint: %w", err) } - key := fmt.Sprintf("checkpoint:%s:%s", checkpoint.ThreadID, checkpoint.ID) + key := redisCheckpointKey(checkpoint.ThreadID, checkpoint.ID) if err := r.client.Set(ctx, key, data, r.ttl).Err(); err != nil { return fmt.Errorf("failed to save checkpoint to Redis: %w", err) } // Add to thread index - threadKey := fmt.Sprintf("thread:%s:checkpoints", checkpoint.ThreadID) + threadKey := redisThreadIndexKey(checkpoint.ThreadID) if err := r.client.SAdd(ctx, threadKey, checkpoint.ID).Err(); err != nil { return fmt.Errorf("failed to add checkpoint to thread index: %w", err) } + // The index set was created without an expiry while the checkpoints it + // points at expire, so it accumulated dead member IDs forever -- an + // unbounded leak that also made List do a wasted round trip per dead entry. + // Refreshing it alongside the newest checkpoint keeps the two in step. + if r.ttl > 0 { + if err := r.client.Expire(ctx, threadKey, r.ttl).Err(); err != nil { + return fmt.Errorf("failed to set thread index expiry: %w", err) + } + } + return nil } // Load loads a checkpoint from Redis func (r *RedisCheckpointer) Load(ctx context.Context, threadID, checkpointID string) (*Checkpoint, error) { - key := fmt.Sprintf("checkpoint:%s:%s", threadID, checkpointID) + key := redisCheckpointKey(threadID, checkpointID) data, err := r.client.Get(ctx, key).Result() if err != nil { @@ -710,22 +1079,33 @@ func (r *RedisCheckpointer) Load(ctx context.Context, threadID, checkpointID str return nil, fmt.Errorf("failed to unmarshal checkpoint: %w", err) } + // Defense in depth against key aliasing: the stored payload records which + // thread it belongs to, so refuse to hand a caller another thread's state + // even if some future key scheme lets two identifiers map to one key. + if checkpoint.ThreadID != "" && checkpoint.ThreadID != threadID { + return nil, fmt.Errorf("checkpoint %s belongs to thread %s, not %s", checkpointID, checkpoint.ThreadID, threadID) + } + return &checkpoint, nil } // List lists checkpoints for a thread func (r *RedisCheckpointer) List(ctx context.Context, threadID string) ([]*CheckpointMetadata, error) { - threadKey := fmt.Sprintf("thread:%s:checkpoints", threadID) + threadKey := redisThreadIndexKey(threadID) checkpointIDs, err := r.client.SMembers(ctx, threadKey).Result() if err != nil { return nil, fmt.Errorf("failed to get checkpoint IDs: %w", err) } - var metadata []*CheckpointMetadata + metadata := []*CheckpointMetadata{} for _, checkpointID := range checkpointIDs { checkpoint, err := r.Load(ctx, threadID, checkpointID) if err != nil { + // An index entry whose checkpoint has expired or been corrupted is + // skipped rather than failing the whole listing, so one bad entry + // cannot make a thread unresumable. It is logged because a silent + // skip would hide real data loss. r.logger.Warnf("Failed to load checkpoint %s: %v", checkpointID, err) continue } @@ -746,18 +1126,26 @@ func (r *RedisCheckpointer) List(ctx context.Context, threadID string) ([]*Check // Delete deletes a checkpoint func (r *RedisCheckpointer) Delete(ctx context.Context, threadID, checkpointID string) error { - key := fmt.Sprintf("checkpoint:%s:%s", threadID, checkpointID) + key := redisCheckpointKey(threadID, checkpointID) - if err := r.client.Del(ctx, key).Err(); err != nil { + removed, err := r.client.Del(ctx, key).Result() + if err != nil { return fmt.Errorf("failed to delete checkpoint from Redis: %w", err) } - // Remove from thread index - threadKey := fmt.Sprintf("thread:%s:checkpoints", threadID) + // Remove from thread index. This runs even when the payload was already + // gone so an expired checkpoint's index entry still gets cleaned up. + threadKey := redisThreadIndexKey(threadID) if err := r.client.SRem(ctx, threadKey, checkpointID).Err(); err != nil { return fmt.Errorf("failed to remove checkpoint from thread index: %w", err) } + // Matches the memory and file checkpointers, which both report a missing + // checkpoint rather than pretending the delete succeeded. + if removed == 0 { + return fmt.Errorf("checkpoint %s not found in thread %s", checkpointID, threadID) + } + return nil } @@ -840,15 +1228,23 @@ func (sm *SessionManager) GetSession(ctx context.Context, sessionID string) (*Se WHERE id = $1 ` - row := sm.conn.QueryRow(ctx, query, sessionID).(*sql.Row) + if sm.conn == nil { + return nil, fmt.Errorf("session manager has no database connection") + } + row, err := asSQLRow(sm.conn.QueryRow(ctx, query, sessionID)) + if err != nil { + return nil, err + } var session Session var metadataData []byte + // user_id is nullable; scanning NULL straight into a string fails. + var userID sql.NullString - err := row.Scan( + err = row.Scan( &session.ID, &session.ThreadID, - &session.UserID, + &userID, &metadataData, &session.CreatedAt, &session.ExpiresAt, @@ -861,8 +1257,10 @@ func (sm *SessionManager) GetSession(ctx context.Context, sessionID string) (*Se return nil, fmt.Errorf("failed to get session: %w", err) } + session.UserID = userID.String + // Unmarshal metadata - if err := json.Unmarshal(metadataData, &session.Metadata); err != nil { + if err := decodeJSONMap(metadataData, &session.Metadata); err != nil { return nil, fmt.Errorf("failed to unmarshal metadata: %w", err) } @@ -898,14 +1296,24 @@ func (sm *SessionManager) GetThread(ctx context.Context, threadID string) (*Thre WHERE id = $1 ` - row := sm.conn.QueryRow(ctx, query, threadID).(*sql.Row) + if sm.conn == nil { + return nil, fmt.Errorf("session manager has no database connection") + } + row, err := asSQLRow(sm.conn.QueryRow(ctx, query, threadID)) + if err != nil { + return nil, err + } var thread Thread var metadataData []byte + // name is nullable, and threads created implicitly by Save have no name at + // all -- scanning that NULL into a string failed with "converting NULL to + // string is unsupported", making every auto-registered thread unreadable. + var name sql.NullString - err := row.Scan( + err = row.Scan( &thread.ID, - &thread.Name, + &name, &metadataData, &thread.CreatedAt, &thread.UpdatedAt, @@ -918,8 +1326,10 @@ func (sm *SessionManager) GetThread(ctx context.Context, threadID string) (*Thre return nil, fmt.Errorf("failed to get thread: %w", err) } + thread.Name = name.String + // Unmarshal metadata - if err := json.Unmarshal(metadataData, &thread.Metadata); err != nil { + if err := decodeJSONMap(metadataData, &thread.Metadata); err != nil { return nil, fmt.Errorf("failed to unmarshal metadata: %w", err) } @@ -928,6 +1338,10 @@ func (sm *SessionManager) GetThread(ctx context.Context, threadID string) (*Thre // DatabaseConnectionManager manages multiple database connections type DatabaseConnectionManager struct { + // mu guards connections. Without it, concurrent AddConnection/GetConnection + // calls are a concurrent map read and write, which the Go runtime turns + // into an unrecoverable process crash rather than a recoverable error. + mu sync.RWMutex connections map[string]DatabaseConnection logger *logrus.Logger } @@ -950,10 +1364,10 @@ func (dcm *DatabaseConnectionManager) AddConnection(name string, config *Databas conn, err = NewPostgresConnection(config) case DatabaseTypeRedis: // Redis connection would be implemented here - return fmt.Errorf("Redis connection not implemented in this version") + return fmt.Errorf("redis connection is not implemented in this version") case DatabaseTypeOpenSearch, DatabaseTypeElastic: // OpenSearch/Elasticsearch connections would be implemented here - return fmt.Errorf("OpenSearch/Elasticsearch connection not implemented in this version") + return fmt.Errorf("openSearch/Elasticsearch connection is not implemented in this version") case DatabaseTypeMongoDB: // MongoDB connection would be implemented here return fmt.Errorf("MongoDB connection not implemented in this version") @@ -971,24 +1385,48 @@ func (dcm *DatabaseConnectionManager) AddConnection(name string, config *Databas return fmt.Errorf("failed to create connection for %s: %w", name, err) } + dcm.mu.Lock() + if dcm.connections == nil { + dcm.connections = make(map[string]DatabaseConnection) + } + // Reusing a name used to overwrite the entry and leak the previous pool, + // which stayed open with no remaining reference for CloseAll to find. + previous, replaced := dcm.connections[name] dcm.connections[name] = conn + dcm.mu.Unlock() + + if replaced && previous != nil { + if cerr := previous.Close(); cerr != nil { + dcm.logger.Warnf("Failed to close replaced connection %s: %v", name, cerr) + } + } + dcm.logger.Infof("Added database connection: %s (%s)", name, config.Type) return nil } // GetConnection retrieves a database connection func (dcm *DatabaseConnectionManager) GetConnection(name string) (DatabaseConnection, error) { + dcm.mu.RLock() conn, exists := dcm.connections[name] + dcm.mu.RUnlock() + if !exists { return nil, fmt.Errorf("connection %s not found", name) } return conn, nil } -// CloseAll closes all database connections +// CloseAll closes all database connections and forgets them, so a second call +// cannot double-close a pool. func (dcm *DatabaseConnectionManager) CloseAll() error { + dcm.mu.Lock() + conns := dcm.connections + dcm.connections = make(map[string]DatabaseConnection) + dcm.mu.Unlock() + var errors []string - for name, conn := range dcm.connections { + for name, conn := range conns { if err := conn.Close(); err != nil { errors = append(errors, fmt.Sprintf("failed to close %s: %v", name, err)) } diff --git a/pkg/persistence/database_driver_test.go b/pkg/persistence/database_driver_test.go new file mode 100644 index 0000000..99e09b8 --- /dev/null +++ b/pkg/persistence/database_driver_test.go @@ -0,0 +1,526 @@ +// Copyright (c) 2024 GoLangGraph Team +// +// Licensed under the MIT License. See LICENSE file in the project root for full license information. +// +// Package: GoLangGraph - A powerful Go framework for building AI agent workflows + +// Driver-level tests for PostgresCheckpointer. +// +// IMPORTANT: the type registered here is a database/sql DRIVER DOUBLE, not a +// real PostgreSQL server. It speaks the driver.Driver/Conn/Stmt/Rows protocol, +// so the checkpointer's own query construction, argument binding, row scanning +// and error handling all run for real -- but no SQL is parsed or executed. +// Behavioral coverage against a genuine server lives in +// postgres_integration_test.go; this file exists for the failure modes a real +// server will not produce on demand: +// +// - a result set that fails partway through iteration, which is what +// rows.Err() is for; +// - a RowsAffected() that reports an error; +// - the exact SQL text and bound parameters the checkpointer sends. +// +// These tests need no database and always run. + +package persistence + +import ( + "context" + "database/sql" + "database/sql/driver" + "errors" + "fmt" + "io" + "sync" + "testing" + "time" + + "github.com/sirupsen/logrus" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +const scriptDriverName = "golanggraph_script_test" + +func init() { + sql.Register(scriptDriverName, scriptDriver{}) +} + +// recordedCall captures one statement the checkpointer sent. +type recordedCall struct { + query string + args []driver.NamedValue +} + +// script tells the double how to answer. One script is registered per test +// under a unique DSN, which is how a connection finds its instructions. +type script struct { + mu sync.Mutex + + columns []string + rows [][]driver.Value + + // failAfter rows have been handed over, Next returns failErr instead of + // io.EOF. A negative value means the result set ends cleanly. + failAfter int + failErr error + + queryErr error + + execAffected int64 + execErr error // returned by RowsAffected, not by Exec itself + + calls []recordedCall +} + +func (s *script) record(query string, args []driver.NamedValue) { + s.mu.Lock() + defer s.mu.Unlock() + s.calls = append(s.calls, recordedCall{query: query, args: args}) +} + +func (s *script) recorded() []recordedCall { + s.mu.Lock() + defer s.mu.Unlock() + out := make([]recordedCall, len(s.calls)) + copy(out, s.calls) + return out +} + +var ( + scriptsMu sync.Mutex + scripts = map[string]*script{} + scriptSeq int +) + +// newScriptedCheckpointer wires a real PostgresCheckpointer on top of the +// double. It deliberately bypasses NewPostgresCheckpointer, whose Connect hard- +// codes the "postgres" driver name and would also try to create the schema. +// Everything below the connection -- the checkpointer's own logic -- is real. +func newScriptedCheckpointer(t *testing.T, s *script, cfg *DatabaseConfig) *PostgresCheckpointer { + t.Helper() + + scriptsMu.Lock() + scriptSeq++ + dsn := fmt.Sprintf("script-%d", scriptSeq) + scripts[dsn] = s + scriptsMu.Unlock() + + t.Cleanup(func() { + scriptsMu.Lock() + delete(scripts, dsn) + scriptsMu.Unlock() + }) + + db, err := sql.Open(scriptDriverName, dsn) + require.NoError(t, err) + t.Cleanup(func() { _ = db.Close() }) + + if cfg == nil { + cfg = NewPostgresConfig("fake", 5432, "fake", "fake", "fake") + } + + return &PostgresCheckpointer{ + conn: &PostgresConnection{db: db, config: cfg, logger: logrus.New()}, + config: cfg, + logger: logrus.New(), + } +} + +// --- the double ----------------------------------------------------------- + +type scriptDriver struct{} + +func (scriptDriver) Open(dsn string) (driver.Conn, error) { + scriptsMu.Lock() + s, ok := scripts[dsn] + scriptsMu.Unlock() + if !ok { + return nil, fmt.Errorf("no script registered for %q", dsn) + } + return &scriptConn{s: s}, nil +} + +type scriptConn struct{ s *script } + +// Implementing QueryerContext and ExecerContext lets database/sql hand us the +// statement and its bound arguments directly, which is what makes the +// assertions on generated SQL possible. +var ( + _ driver.QueryerContext = (*scriptConn)(nil) + _ driver.ExecerContext = (*scriptConn)(nil) +) + +func (c *scriptConn) QueryContext(ctx context.Context, query string, args []driver.NamedValue) (driver.Rows, error) { + c.s.record(query, args) + if c.s.queryErr != nil { + return nil, c.s.queryErr + } + return &scriptRows{s: c.s}, nil +} + +func (c *scriptConn) ExecContext(ctx context.Context, query string, args []driver.NamedValue) (driver.Result, error) { + c.s.record(query, args) + return scriptResult{affected: c.s.execAffected, err: c.s.execErr}, nil +} + +func (c *scriptConn) Prepare(query string) (driver.Stmt, error) { + return &scriptStmt{c: c, query: query}, nil +} +func (c *scriptConn) Close() error { return nil } +func (c *scriptConn) Begin() (driver.Tx, error) { return scriptTx{}, nil } + +type scriptTx struct{} + +func (scriptTx) Commit() error { return nil } +func (scriptTx) Rollback() error { return nil } + +type scriptStmt struct { + c *scriptConn + query string +} + +func (s *scriptStmt) Close() error { return nil } +func (s *scriptStmt) NumInput() int { return -1 } // -1 disables arity checking +func (s *scriptStmt) Exec(args []driver.Value) (driver.Result, error) { + return scriptResult{affected: s.c.s.execAffected, err: s.c.s.execErr}, nil +} +func (s *scriptStmt) Query(args []driver.Value) (driver.Rows, error) { + return &scriptRows{s: s.c.s}, nil +} + +type scriptResult struct { + affected int64 + err error +} + +func (r scriptResult) LastInsertId() (int64, error) { return 0, nil } +func (r scriptResult) RowsAffected() (int64, error) { return r.affected, r.err } + +type scriptRows struct { + s *script + pos int +} + +func (r *scriptRows) Columns() []string { return r.s.columns } +func (r *scriptRows) Close() error { return nil } + +func (r *scriptRows) Next(dest []driver.Value) error { + if r.s.failAfter >= 0 && r.pos >= r.s.failAfter { + // A non-EOF error here is exactly what a connection dropping mid-result + // looks like to database/sql: Rows.Next reports false and the reason is + // only available from Rows.Err. + return r.s.failErr + } + if r.pos >= len(r.s.rows) { + return io.EOF + } + copy(dest, r.s.rows[r.pos]) + r.pos++ + return nil +} + +// checkpointRow builds one row shaped like the List query's SELECT list: +// id, thread_id, metadata, created_at, node_id, step_id. +func checkpointRow(id, threadID string, nodeID interface{}, stepID interface{}) []driver.Value { + return []driver.Value{ + id, + threadID, + []byte(`{"k":"v"}`), + time.Date(2024, 1, 1, 12, 0, 0, 0, time.UTC), + nodeID, + stepID, + } +} + +var checkpointColumns = []string{"id", "thread_id", "metadata", "created_at", "node_id", "step_id"} + +// --- rows.Err() ----------------------------------------------------------- + +// TestPostgresCheckpointer_ListSurfacesIterationError covers a silent +// truncation bug: List iterated with `for rows.Next()` and returned without ever +// calling rows.Err(). A connection that dropped after two of five rows produced +// a two-element slice and a nil error, so the caller could not tell a partial +// answer from a complete one -- and Latest() would happily "resume" from a +// checkpoint that was not actually the newest. +func TestPostgresCheckpointer_ListSurfacesIterationError(t *testing.T) { + boom := errors.New("connection reset by peer") + s := &script{ + columns: checkpointColumns, + rows: [][]driver.Value{ + checkpointRow("ckpt-1", "t1", "node-1", int64(1)), + checkpointRow("ckpt-2", "t1", "node-2", int64(2)), + checkpointRow("ckpt-3", "t1", "node-3", int64(3)), + }, + failAfter: 2, // two rows arrive, then the connection dies + failErr: boom, + } + + cp := newScriptedCheckpointer(t, s, nil) + + got, err := cp.List(context.Background(), "t1") + require.Error(t, err, "a result set that fails mid-iteration must not look like a complete listing") + assert.ErrorIs(t, err, boom) + assert.Nil(t, got, "no partial slice may be handed back alongside the error") +} + +// TestPostgresCheckpointer_SearchDocumentsSurfacesIterationError is the same +// unchecked-rows.Err() defect in the RAG search path. +func TestPostgresCheckpointer_SearchDocumentsSurfacesIterationError(t *testing.T) { + boom := errors.New("server closed the connection unexpectedly") + s := &script{ + columns: []string{"id", "thread_id", "content", "metadata", "created_at", "updated_at"}, + rows: [][]driver.Value{ + {"doc-1", "t1", "content one", []byte(`{}`), time.Now().UTC(), time.Now().UTC()}, + {"doc-2", "t1", "content two", []byte(`{}`), time.Now().UTC(), time.Now().UTC()}, + }, + failAfter: 1, + failErr: boom, + } + + cfg := NewPostgresConfig("fake", 5432, "fake", "fake", "fake") + cfg.EnableRAG = true + cp := newScriptedCheckpointer(t, s, cfg) + + got, err := cp.SearchDocuments(context.Background(), "t1", nil, 10) + require.Error(t, err, "a truncated document search must be reported") + assert.ErrorIs(t, err, boom) + assert.Nil(t, got) +} + +// TestPostgresCheckpointer_ListSucceedsOnCleanIteration is the control for the +// two tests above: with no injected failure the same code path must return +// every row and no error, so the checks are not just failing on principle. +func TestPostgresCheckpointer_ListSucceedsOnCleanIteration(t *testing.T) { + s := &script{ + columns: checkpointColumns, + rows: [][]driver.Value{ + checkpointRow("ckpt-1", "t1", "node-1", int64(1)), + // Second row exercises NULL node_id/step_id at the driver level. + checkpointRow("ckpt-2", "t1", nil, nil), + }, + failAfter: -1, + } + + cp := newScriptedCheckpointer(t, s, nil) + + got, err := cp.List(context.Background(), "t1") + require.NoError(t, err) + require.Len(t, got, 2) + + assert.Equal(t, "ckpt-1", got[0].ID) + assert.Equal(t, "node-1", got[0].NodeID) + assert.Equal(t, 1, got[0].StepID) + assert.Equal(t, "v", got[0].Metadata["k"], "metadata must be decoded from the raw column bytes") + + assert.Equal(t, "", got[1].NodeID, "a NULL node_id must scan to the empty string") + assert.Equal(t, 0, got[1].StepID, "a NULL step_id must scan to zero") +} + +// --- generated SQL and bound parameters ----------------------------------- + +// TestPostgresCheckpointer_ListBindsThreadIDAsParameter pins the shape of the +// generated statement. The thread ID must travel as a bound parameter rather +// than being concatenated into the SQL text, which is what keeps a hostile +// thread ID inert. +func TestPostgresCheckpointer_ListBindsThreadIDAsParameter(t *testing.T) { + s := &script{columns: checkpointColumns, failAfter: -1} + cp := newScriptedCheckpointer(t, s, nil) + + hostile := `t1'; DROP TABLE checkpoints; --` + _, err := cp.List(context.Background(), hostile) + require.NoError(t, err) + + calls := s.recorded() + require.Len(t, calls, 1) + + assert.Contains(t, calls[0].query, "FROM checkpoints") + assert.Contains(t, calls[0].query, "WHERE thread_id = $1") + assert.Contains(t, calls[0].query, "ORDER BY created_at DESC") + assert.NotContains(t, calls[0].query, "DROP TABLE", + "the thread ID must never be interpolated into the SQL text") + + require.Len(t, calls[0].args, 1) + assert.Equal(t, hostile, calls[0].args[0].Value, "the thread ID must arrive as a bound parameter") +} + +// TestPostgresCheckpointer_DeleteBindsBothIdentifiers checks the delete is +// scoped by thread as well as by checkpoint ID, so one thread cannot delete +// another's checkpoint by guessing its ID. +func TestPostgresCheckpointer_DeleteBindsBothIdentifiers(t *testing.T) { + s := &script{execAffected: 1} + cp := newScriptedCheckpointer(t, s, nil) + + require.NoError(t, cp.Delete(context.Background(), "thread-9", "ckpt-9")) + + calls := s.recorded() + require.Len(t, calls, 1) + assert.Contains(t, calls[0].query, "DELETE FROM checkpoints") + assert.Contains(t, calls[0].query, "thread_id = $1") + assert.Contains(t, calls[0].query, "id = $2") + + require.Len(t, calls[0].args, 2) + assert.Equal(t, "thread-9", calls[0].args[0].Value) + assert.Equal(t, "ckpt-9", calls[0].args[1].Value) +} + +// TestPostgresCheckpointer_SaveRegistersThreadBeforeCheckpoint pins the fix for +// the foreign key violation: the parent thread row must be written first, and +// both statements must go out together so a failed checkpoint cannot leave an +// orphan thread behind. +func TestPostgresCheckpointer_SaveRegistersThreadBeforeCheckpoint(t *testing.T) { + s := &script{execAffected: 1} + cp := newScriptedCheckpointer(t, s, nil) + + ck := newCheckpoint("thread-x", "ckpt-x", 2, richState()) + require.NoError(t, cp.Save(context.Background(), ck)) + + calls := s.recorded() + require.Len(t, calls, 2, "Save must issue the thread upsert and the checkpoint upsert") + + assert.Contains(t, calls[0].query, "INSERT INTO threads") + assert.Contains(t, calls[0].query, "ON CONFLICT (id) DO NOTHING", + "registering an existing thread must not be an error") + require.Len(t, calls[0].args, 1) + assert.Equal(t, "thread-x", calls[0].args[0].Value) + + assert.Contains(t, calls[1].query, "INSERT INTO checkpoints") + assert.Contains(t, calls[1].query, "ON CONFLICT (id) DO UPDATE") + require.Len(t, calls[1].args, 7) + assert.Equal(t, "ckpt-x", calls[1].args[0].Value) + assert.Equal(t, "thread-x", calls[1].args[1].Value) + + // The state must reach the driver as real JSON. This is the same regression + // the integration tests cover, asserted one layer lower: before BaseState + // implemented json.Marshaler this argument was the two bytes "{}". + stateArg, ok := calls[1].args[2].Value.([]byte) + require.True(t, ok, "state should be bound as raw JSON bytes, got %T", calls[1].args[2].Value) + assert.NotEqual(t, "{}", string(stateArg), "state was serialized as an empty object -- all state lost") + assert.Contains(t, string(stateArg), "hello world") +} + +// --- error handling in results -------------------------------------------- + +// TestPostgresCheckpointer_DeleteReportsRowsAffectedError covers the branch +// where the driver cannot say how many rows were removed. Reporting success +// there would claim a deletion that may not have happened. +func TestPostgresCheckpointer_DeleteReportsRowsAffectedError(t *testing.T) { + boom := errors.New("no RowsAffected available") + s := &script{execErr: boom} + cp := newScriptedCheckpointer(t, s, nil) + + err := cp.Delete(context.Background(), "t", "c") + require.Error(t, err) + assert.ErrorIs(t, err, boom) +} + +// TestPostgresCheckpointer_DeleteMissingRowIsNotFound is the driver-level twin +// of the integration test: zero rows affected must be an error. +func TestPostgresCheckpointer_DeleteMissingRowIsNotFound(t *testing.T) { + s := &script{execAffected: 0} + cp := newScriptedCheckpointer(t, s, nil) + + err := cp.Delete(context.Background(), "t", "c") + require.Error(t, err) + assert.Contains(t, err.Error(), "not found") +} + +func TestPostgresCheckpointer_ListReportsQueryError(t *testing.T) { + boom := errors.New("relation \"checkpoints\" does not exist") + s := &script{columns: checkpointColumns, failAfter: -1, queryErr: boom} + cp := newScriptedCheckpointer(t, s, nil) + + _, err := cp.List(context.Background(), "t1") + require.Error(t, err) + assert.ErrorIs(t, err, boom) +} + +// --- guards against non-SQL connections ----------------------------------- + +// TestAsSQLRows_RejectsForeignTypes covers the helper that replaced inline +// `rows.(*sql.Rows)` assertions. Those are single-value assertions, so any +// DatabaseConnection implementation other than PostgresConnection -- the +// interface returns interface{}, so others are allowed -- panicked and took the +// process down instead of returning an error. +func TestAsSQLRows_RejectsForeignTypes(t *testing.T) { + _, err := asSQLRows(nil) + assert.Error(t, err, "nil must be an error, not a panic") + + _, err = asSQLRows("not rows") + require.Error(t, err) + assert.Contains(t, err.Error(), "want *sql.Rows") + + var typedNil *sql.Rows + _, err = asSQLRows(typedNil) + assert.Error(t, err, "a typed nil must be rejected too, or the caller nil-derefs") +} + +func TestAsSQLRow_RejectsForeignTypes(t *testing.T) { + _, err := asSQLRow(nil) + assert.Error(t, err) + + _, err = asSQLRow(struct{}{}) + require.Error(t, err) + assert.Contains(t, err.Error(), "want *sql.Row") +} + +// --- pure helpers --------------------------------------------------------- + +func TestEncodeDecodeVector(t *testing.T) { + cases := [][]float64{ + {1, 2, 3}, + {-0.5, 0.25, 1e-7}, + {0}, + } + for _, want := range cases { + encoded := encodeVector(want) + got, err := decodeVector(encoded) + require.NoError(t, err) + assert.Equal(t, want, got, "vector must survive encode/decode exactly") + } + + assert.Equal(t, "[]", encodeVector(nil)) + + // SQL NULL must stay absent rather than becoming an empty slice, so callers + // can tell "no embedding stored" from "a zero-length embedding". + got, err := decodeVector(nil) + require.NoError(t, err) + assert.Nil(t, got) + + got, err = decodeVector([]byte("[1,2]")) + require.NoError(t, err) + assert.Equal(t, []float64{1, 2}, got, "a []byte column value must decode as well as a string") + + _, err = decodeVector("[1,notanumber]") + assert.Error(t, err, "a malformed vector must be reported, not silently truncated") + + _, err = decodeVector(42) + assert.Error(t, err, "an unexpected column type must be reported") +} + +func TestDecodeJSONMap(t *testing.T) { + var m map[string]interface{} + + // SQL NULL arrives as a nil slice; the JSON literal null is also possible. + require.NoError(t, decodeJSONMap(nil, &m)) + assert.NotNil(t, m, "NULL must decode to an empty map, never nil -- callers write into this map") + assert.Empty(t, m) + + require.NoError(t, decodeJSONMap([]byte("null"), &m)) + assert.NotNil(t, m) + + require.NoError(t, decodeJSONMap([]byte(`{"a":1}`), &m)) + assert.EqualValues(t, 1, m["a"]) + + assert.Error(t, decodeJSONMap([]byte("{broken"), &m)) +} + +// TestDecodeJSONMap_ResultIsWritable guards the reason NULL must not decode to +// nil: writing to a nil map panics, and callers treat checkpoint metadata as a +// normal map. +func TestDecodeJSONMap_ResultIsWritable(t *testing.T) { + var m map[string]interface{} + require.NoError(t, decodeJSONMap(nil, &m)) + + assert.NotPanics(t, func() { m["added"] = true }) + assert.Equal(t, true, m["added"]) +} diff --git a/pkg/persistence/database_test.go b/pkg/persistence/database_test.go index 69c9c37..3bf5190 100644 --- a/pkg/persistence/database_test.go +++ b/pkg/persistence/database_test.go @@ -90,13 +90,20 @@ func TestNewDatabaseConnectionManager(t *testing.T) { func TestDatabaseConnectionManager_AddConnection(t *testing.T) { manager := NewDatabaseConnectionManager() - // Test adding PostgreSQL connection - config := NewPostgresConfig("localhost", 5432, "testdb", "testuser", "testpass") + // Port 1 is reserved and never listening, so this exercises the failure + // path deterministically. Using a real port such as 5432 would make the + // test depend on no PostgreSQL running locally, which is false on most + // developer machines and in any CI job with a database service. + config := NewPostgresConfig("127.0.0.1", 1, "testdb", "testuser", "testpass") err := manager.AddConnection("test_postgres", config) - // This will fail without actual database, but we can test the error handling if err == nil { - t.Error("AddConnection should return error when database is not available") + t.Error("AddConnection should return error when the database is unreachable") + } + + // A failed AddConnection must not leave a half-built entry behind. + if _, getErr := manager.GetConnection("test_postgres"); getErr == nil { + t.Error("A connection that failed to open should not be registered") } } @@ -111,22 +118,21 @@ func TestDatabaseConnectionManager_GetConnection(t *testing.T) { } func TestCreateCheckpointer(t *testing.T) { - // Test creating PostgreSQL checkpointer - config := NewPostgresConfig("localhost", 5432, "testdb", "testuser", "testpass") + // Port 1 is reserved and never listening, so the unreachable-server branch + // is exercised deterministically. Pointing at the real default ports would + // make these assertions fail wherever a database actually runs. + config := NewPostgresConfig("127.0.0.1", 1, "testdb", "testuser", "testpass") _, err := CreateCheckpointer(config) - // This will fail without actual database, but we can test the error handling if err == nil { - t.Error("CreateCheckpointer should return error when database is not available") + t.Error("CreateCheckpointer should return error when the database is unreachable") } - // Test creating Redis checkpointer - redisConfig := NewRedisConfig("localhost", 6379, "testpass") + redisConfig := NewRedisConfig("127.0.0.1", 1, "testpass") _, err = CreateCheckpointer(redisConfig) - // This will fail without actual Redis, but we can test the error handling if err == nil { - t.Error("CreateCheckpointer should return error when Redis is not available") + t.Error("CreateCheckpointer should return error when Redis is unreachable") } // Test creating checkpointer with unsupported type @@ -324,26 +330,11 @@ func (m *MockConnection) QueryRows(ctx context.Context, query string, args ...in return nil, nil } -// Integration tests (these would require actual database connections) -func TestPostgresCheckpointer_Integration(t *testing.T) { - if testing.Short() { - t.Skip("Skipping integration test in short mode") - } - - // This test would require a real PostgreSQL database - // For now, we'll skip it unless explicitly running integration tests - t.Skip("Integration test requires PostgreSQL database") -} - -func TestRedisCheckpointer_Integration(t *testing.T) { - if testing.Short() { - t.Skip("Skipping integration test in short mode") - } - - // This test would require a real Redis instance - // For now, we'll skip it unless explicitly running integration tests - t.Skip("Integration test requires Redis instance") -} +// The unconditional t.Skip() placeholders that used to sit here never executed +// a single line of PostgresCheckpointer or RedisCheckpointer. Real integration +// coverage now lives in postgres_integration_test.go and +// redis_integration_test.go, which run against actual servers and skip with an +// explanation only when none is reachable. // Benchmark tests func BenchmarkNewPostgresConfig(b *testing.B) { diff --git a/pkg/persistence/persistence_test.go b/pkg/persistence/persistence_test.go new file mode 100644 index 0000000..ef2cc8e --- /dev/null +++ b/pkg/persistence/persistence_test.go @@ -0,0 +1,347 @@ +// Copyright (c) 2024 GoLangGraph Team +// +// Licensed under the MIT License. See LICENSE file in the project root for full license information. + +package persistence + +import ( + "context" + "encoding/json" + "fmt" + "os" + "path/filepath" + "strings" + "sync" + "testing" + "time" + + "github.com/UnicoLab/GoLangGraph/pkg/core" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func stateWith(kv map[string]interface{}) *core.BaseState { + s := core.NewBaseState() + for k, v := range kv { + s.Set(k, v) + } + return s +} + +// A checkpoint must survive JSON encoding. BaseState keeps its data in +// unexported fields, so without a marshaller every persisted checkpoint and +// every API response carrying a state is silently empty. +func TestCheckpoint_JSONRoundTripPreservesState(t *testing.T) { + cp := &Checkpoint{ + ID: "c1", ThreadID: "t1", NodeID: "n1", StepID: 3, + CreatedAt: time.Now().UTC().Truncate(time.Second), + State: stateWith(map[string]interface{}{"counter": 42, "messages": []interface{}{"hi"}}), + Metadata: map[string]interface{}{"source": "test"}, + } + + data, err := json.Marshal(cp) + require.NoError(t, err) + assert.NotContains(t, string(data), `"state":{}`, "state must not serialize as an empty object") + + var back Checkpoint + require.NoError(t, json.Unmarshal(data, &back)) + require.NotNil(t, back.State) + + counter, ok := back.State.Get("counter") + require.True(t, ok, "state was lost across the round trip") + assert.EqualValues(t, 42, counter) + assert.Equal(t, 3, back.StepID) + assert.Equal(t, "n1", back.NodeID) +} + +// FileCheckpointer previously had stub IO: Save reported success and wrote +// nothing at all. +func TestFileCheckpointer_ActuallyWritesToDisk(t *testing.T) { + base := filepath.Join(t.TempDir(), "store") + cp := NewFileCheckpointer(base) + ctx := context.Background() + + require.NoError(t, cp.Save(ctx, &Checkpoint{ + ID: "cp-1", ThreadID: "t1", NodeID: "start", StepID: 0, + State: stateWith(map[string]interface{}{"v": "persisted"}), CreatedAt: time.Now(), + })) + + path := filepath.Join(base, "t1", "cp-1.json") + info, err := os.Stat(path) + require.NoError(t, err, "no checkpoint file was written") + assert.Greater(t, info.Size(), int64(0)) + + raw, err := os.ReadFile(path) + require.NoError(t, err) + assert.Contains(t, string(raw), "persisted") + + loaded, err := cp.Load(ctx, "t1", "cp-1") + require.NoError(t, err) + v, ok := loaded.State.Get("v") + require.True(t, ok) + assert.Equal(t, "persisted", v) +} + +// A checkpoint survives a process restart: a fresh checkpointer over the same +// directory sees everything the previous one wrote. +func TestFileCheckpointer_SurvivesRestart(t *testing.T) { + base := filepath.Join(t.TempDir(), "store") + ctx := context.Background() + + first := NewFileCheckpointer(base) + for i := 0; i < 3; i++ { + require.NoError(t, first.Save(ctx, &Checkpoint{ + ID: fmt.Sprintf("cp-%d", i), ThreadID: "t", StepID: i, + State: stateWith(map[string]interface{}{"step": i}), CreatedAt: time.Now(), + })) + } + require.NoError(t, first.Close()) + + // A new process attaches to the same directory. + second := NewFileCheckpointer(base) + metas, err := second.List(ctx, "t") + require.NoError(t, err) + assert.Len(t, metas, 3, "checkpoints must outlive the process that wrote them") + + latest, err := Latest(ctx, second, "t") + require.NoError(t, err) + require.NotNil(t, latest) + step, _ := latest.State.Get("step") + assert.EqualValues(t, 2, step) +} + +// Writes are atomic: a reader must never observe a half-written checkpoint. +func TestFileCheckpointer_NoPartialFilesRemain(t *testing.T) { + base := filepath.Join(t.TempDir(), "store") + cp := NewFileCheckpointer(base) + ctx := context.Background() + + big := strings.Repeat("payload", 10000) + require.NoError(t, cp.Save(ctx, &Checkpoint{ + ID: "cp", ThreadID: "t", State: stateWith(map[string]interface{}{"blob": big}), CreatedAt: time.Now(), + })) + + entries, err := os.ReadDir(filepath.Join(base, "t")) + require.NoError(t, err) + for _, e := range entries { + assert.False(t, strings.HasSuffix(e.Name(), ".tmp"), + "temporary file %s was left behind", e.Name()) + } + assert.Len(t, entries, 1) +} + +func TestFileCheckpointer_RejectsUnsafeIdentifiers(t *testing.T) { + cp := NewFileCheckpointer(filepath.Join(t.TempDir(), "store")) + ctx := context.Background() + + for _, bad := range []string{"../escape", "a/b", "..", "", "x\x00y", strings.Repeat("a", 300)} { + err := cp.Save(ctx, &Checkpoint{ID: "ok", ThreadID: bad, State: core.NewBaseState(), CreatedAt: time.Now()}) + assert.Error(t, err, "thread ID %q must be rejected", bad) + + _, err = cp.Load(ctx, bad, "ok") + assert.Error(t, err, "loading thread ID %q must be rejected", bad) + } +} + +func TestMemoryCheckpointer_IsolatesStoredState(t *testing.T) { + cp := NewMemoryCheckpointer() + ctx := context.Background() + + original := stateWith(map[string]interface{}{"v": 1}) + require.NoError(t, cp.Save(ctx, &Checkpoint{ID: "cp", ThreadID: "t", State: original, CreatedAt: time.Now()})) + + // Mutating the caller's state must not change what was stored. + original.Set("v", 999) + + loaded, err := cp.Load(ctx, "t", "cp") + require.NoError(t, err) + v, _ := loaded.State.Get("v") + assert.EqualValues(t, 1, v, "stored checkpoints must not alias caller state") + + // Mutating a loaded copy must not change the store either. + loaded.State.Set("v", 555) + again, err := cp.Load(ctx, "t", "cp") + require.NoError(t, err) + v, _ = again.State.Get("v") + assert.EqualValues(t, 1, v) +} + +// Every backend must behave the same under concurrent use. +func TestCheckpointers_ConcurrentAccess(t *testing.T) { + backends := map[string]Checkpointer{ + "memory": NewMemoryCheckpointer(), + "file": NewFileCheckpointer(filepath.Join(t.TempDir(), "store")), + } + + for name, cp := range backends { + t.Run(name, func(t *testing.T) { + ctx := context.Background() + var wg sync.WaitGroup + errs := make([]error, 24) + + for i := 0; i < 24; i++ { + wg.Add(1) + go func(i int) { + defer wg.Done() + thread := fmt.Sprintf("t%d", i%4) + id := fmt.Sprintf("cp-%d", i) + if err := cp.Save(ctx, &Checkpoint{ + ID: id, ThreadID: thread, StepID: i, + State: stateWith(map[string]interface{}{"i": i}), CreatedAt: time.Now(), + }); err != nil { + errs[i] = err + return + } + if _, err := cp.Load(ctx, thread, id); err != nil { + errs[i] = err + return + } + if _, err := cp.List(ctx, thread); err != nil { + errs[i] = err + } + }(i) + } + wg.Wait() + for i, err := range errs { + assert.NoError(t, err, "worker %d", i) + } + }) + } +} + +// A corrupt file must be reported on load and skipped on list, never panic. +func TestFileCheckpointer_CorruptedDataIsHandled(t *testing.T) { + base := filepath.Join(t.TempDir(), "store") + cp := NewFileCheckpointer(base) + ctx := context.Background() + + require.NoError(t, cp.Save(ctx, &Checkpoint{ + ID: "good", ThreadID: "t", State: stateWith(map[string]interface{}{"v": 1}), CreatedAt: time.Now(), + })) + require.NoError(t, cp.Save(ctx, &Checkpoint{ + ID: "bad", ThreadID: "t", State: stateWith(map[string]interface{}{"v": 2}), CreatedAt: time.Now(), + })) + + require.NoError(t, os.WriteFile(filepath.Join(base, "t", "bad.json"), []byte("\x00\x01 not json"), 0o600)) + + _, err := cp.Load(ctx, "t", "bad") + assert.Error(t, err, "a corrupt checkpoint must be reported") + + metas, err := cp.List(ctx, "t") + require.NoError(t, err, "one corrupt file must not fail the whole listing") + require.Len(t, metas, 1) + assert.Equal(t, "good", metas[0].ID) +} + +func TestCheckpointSaver_PersistsPerStep(t *testing.T) { + cp := NewMemoryCheckpointer() + saver := NewCheckpointSaver(cp) + ctx := context.Background() + + for step, node := range []string{"a", "b", "c"} { + require.NoError(t, saver.SaveState(ctx, "thread", node, step, + stateWith(map[string]interface{}{"step": step}))) + } + + metas, err := cp.List(ctx, "thread") + require.NoError(t, err) + assert.Len(t, metas, 3) + + latest, err := Latest(ctx, cp, "thread") + require.NoError(t, err) + assert.Equal(t, "c", latest.NodeID) + assert.Equal(t, 2, latest.StepID) +} + +func TestCheckpointSaver_RejectsNilState(t *testing.T) { + saver := NewCheckpointSaver(NewMemoryCheckpointer()) + err := saver.SaveState(context.Background(), "t", "n", 0, nil) + assert.Error(t, err) +} + +func TestCheckpointSaver_NilCheckpointerIsNoOp(t *testing.T) { + saver := NewCheckpointSaver(nil) + assert.NoError(t, saver.SaveState(context.Background(), "t", "n", 0, core.NewBaseState())) +} + +func TestLatest_EmptyThread(t *testing.T) { + latest, err := Latest(context.Background(), NewMemoryCheckpointer(), "nothing-here") + require.NoError(t, err) + assert.Nil(t, latest, "an empty thread has no latest checkpoint, and that is not an error") +} + +// --------------------------------------------------------------------------- +// Database layer robustness +// --------------------------------------------------------------------------- + +// fakeConnection implements DatabaseConnection without a real database, which +// is exactly the case that used to panic on an unchecked type assertion. +type fakeConnection struct { + execErr error + pinged bool +} + +func (f *fakeConnection) Connect() error { return nil } +func (f *fakeConnection) Ping() error { f.pinged = true; return nil } +func (f *fakeConnection) Close() error { return nil } +func (f *fakeConnection) GetType() DatabaseType { return DatabaseType("fake") } +func (f *fakeConnection) GetConfig() *DatabaseConfig { + return &DatabaseConfig{Type: DatabaseType("fake")} +} +func (f *fakeConnection) ExecuteQuery(ctx context.Context, query string, args ...interface{}) error { + return f.execErr +} +func (f *fakeConnection) QueryRow(ctx context.Context, query string, args ...interface{}) interface{} { + return nil +} +func (f *fakeConnection) QueryRows(ctx context.Context, query string, args ...interface{}) (interface{}, error) { + return nil, nil +} + +func TestSessionManager_NonSQLConnectionReturnsErrorNotPanic(t *testing.T) { + sm := NewSessionManager(&fakeConnection{}) + ctx := context.Background() + + defer func() { + if r := recover(); r != nil { + t.Fatalf("a non-SQL DatabaseConnection must not panic: %v", r) + } + }() + + _, err := sm.GetSession(ctx, "s1") + assert.Error(t, err) + + _, err = sm.GetThread(ctx, "t1") + assert.Error(t, err) +} + +func TestSessionManager_NilConnectionReturnsError(t *testing.T) { + sm := NewSessionManager(nil) + ctx := context.Background() + + defer func() { + if r := recover(); r != nil { + t.Fatalf("a nil connection must not panic: %v", r) + } + }() + + _, err := sm.GetSession(ctx, "s1") + assert.Error(t, err) +} + +func TestPostgresConnection_UnopenedConnectionErrors(t *testing.T) { + conn := &PostgresConnection{} + ctx := context.Background() + + defer func() { + if r := recover(); r != nil { + t.Fatalf("an unopened connection must not panic: %v", r) + } + }() + + assert.Error(t, conn.Ping()) + assert.Error(t, conn.ExecuteQuery(ctx, "SELECT 1")) + _, err := conn.QueryRows(ctx, "SELECT 1") + assert.Error(t, err) + assert.Nil(t, conn.QueryRow(ctx, "SELECT 1")) +} diff --git a/pkg/persistence/postgres_integration_test.go b/pkg/persistence/postgres_integration_test.go new file mode 100644 index 0000000..eca4f59 --- /dev/null +++ b/pkg/persistence/postgres_integration_test.go @@ -0,0 +1,940 @@ +// Copyright (c) 2024 GoLangGraph Team +// +// Licensed under the MIT License. See LICENSE file in the project root for full license information. +// +// Package: GoLangGraph - A powerful Go framework for building AI agent workflows + +// Integration tests that run PostgresCheckpointer against a REAL PostgreSQL +// server. Nothing here is mocked: the queries, the argument binding, the row +// scanning and the schema are all exercised end to end. +// +// Discovery order: +// 1. POSTGRES_TEST_DSN, if set. An explicitly configured server that cannot be +// reached is a hard failure -- the operator asked for these tests to run. +// 2. Otherwise a couple of conventional local DSNs. If none answer, every test +// in this file skips with an explanation, so CI without a database stays +// green. +// +// Each test runs inside its own PostgreSQL schema so tests cannot see one +// another's rows and a failure leaves nothing behind. + +package persistence + +import ( + "context" + "database/sql" + "fmt" + "net/url" + "os" + "strconv" + "strings" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/lib/pq" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/UnicoLab/GoLangGraph/pkg/core" +) + +// candidatePostgresDSNs are tried in order when POSTGRES_TEST_DSN is unset. +// They cover the two usual local setups: password auth with the conventional +// "postgres" password, and trust auth where no password is needed. +var candidatePostgresDSNs = []string{ + "postgres://postgres:postgres@127.0.0.1:5432/postgres?sslmode=disable", + "postgres://postgres@127.0.0.1:5432/postgres?sslmode=disable", +} + +var ( + pgDiscoverOnce sync.Once + pgBaseConfig *DatabaseConfig + pgDiscoverErr error + pgExplicitDSN bool + pgHasVector bool +) + +// parsePostgresDSN turns a postgres:// URL into the DatabaseConfig the +// checkpointer expects. The checkpointer builds its own libpq DSN from these +// fields, so we cannot simply pass the URL through. +func parsePostgresDSN(dsn string) (*DatabaseConfig, error) { + u, err := url.Parse(dsn) + if err != nil { + return nil, fmt.Errorf("invalid DSN %q: %w", dsn, err) + } + + port := 5432 + if p := u.Port(); p != "" { + port, err = strconv.Atoi(p) + if err != nil { + return nil, fmt.Errorf("invalid port in DSN %q: %w", dsn, err) + } + } + + password, _ := u.User.Password() + cfg := NewPostgresConfig( + u.Hostname(), + port, + strings.TrimPrefix(u.Path, "/"), + u.User.Username(), + password, + ) + if mode := u.Query().Get("sslmode"); mode != "" { + cfg.SSLMode = mode + } + return cfg, nil +} + +// discoverPostgres finds a usable server exactly once per test binary. +func discoverPostgres() (*DatabaseConfig, error) { + pgDiscoverOnce.Do(func() { + dsns := candidatePostgresDSNs + if explicit := os.Getenv("POSTGRES_TEST_DSN"); explicit != "" { + dsns = []string{explicit} + pgExplicitDSN = true + } + + var lastErr error + for _, dsn := range dsns { + cfg, err := parsePostgresDSN(dsn) + if err != nil { + lastErr = err + continue + } + conn, err := NewPostgresConnection(cfg) + if err != nil { + lastErr = err + continue + } + + // Record whether pgvector is usable so the vector tests can skip + // precisely rather than failing on a server without the extension. + var one int + row, rerr := asSQLRow(conn.QueryRow(context.Background(), + `SELECT 1 FROM pg_extension WHERE extname = 'vector'`)) + if rerr == nil && row.Scan(&one) == nil { + pgHasVector = true + } + + _ = conn.Close() + pgBaseConfig = cfg + return + } + pgDiscoverErr = lastErr + }) + return pgBaseConfig, pgDiscoverErr +} + +// requirePostgres returns a base config, skipping the test when no server is +// reachable. An explicitly requested server that is down fails instead. +func requirePostgres(t *testing.T) *DatabaseConfig { + t.Helper() + + cfg, err := discoverPostgres() + if cfg == nil { + if pgExplicitDSN { + t.Fatalf("POSTGRES_TEST_DSN is set but the server is unreachable: %v", err) + } + t.Skipf("no local PostgreSQL reachable (tried %v; last error: %v). "+ + "Set POSTGRES_TEST_DSN to run these integration tests.", candidatePostgresDSNs, err) + } + + clone := *cfg + return &clone +} + +// newPostgresSchema gives a test its own PostgreSQL schema, dropped on cleanup. +// +// Isolation matters here because the checkpointer creates fixed table names +// (threads, checkpoints, documents...) and one test's rows would otherwise show +// up in another's List. It also means each test starts against a genuinely +// empty schema, which is what exercises initSchema. +func newPostgresSchema(t *testing.T, mutate func(*DatabaseConfig)) *DatabaseConfig { + t.Helper() + + base := requirePostgres(t) + admin, err := NewPostgresConnection(base) + require.NoError(t, err, "connect to PostgreSQL") + + schema := fmt.Sprintf("gltest_%d_%d", time.Now().UnixNano(), atomic.AddUint64(&schemaCounter, 1)) + ctx := context.Background() + require.NoError(t, admin.ExecuteQuery(ctx, "CREATE SCHEMA "+pq.QuoteIdentifier(schema))) + + t.Cleanup(func() { + if err := admin.ExecuteQuery(context.Background(), + "DROP SCHEMA "+pq.QuoteIdentifier(schema)+" CASCADE"); err != nil { + t.Logf("failed to drop schema %s: %v", schema, err) + } + if err := admin.Close(); err != nil { + t.Logf("failed to close admin connection: %v", err) + } + }) + + cfg := *base + // "public" stays on the path so extension-provided types such as pgvector's + // "vector" resolve; the test schema comes first so CREATE TABLE lands there. + cfg.ConnectionParams = map[string]string{"search_path": schema + ",public"} + if mutate != nil { + mutate(&cfg) + } + return &cfg +} + +var schemaCounter uint64 + +// newPostgresCheckpointer builds a checkpointer on a private schema. +func newPostgresCheckpointer(t *testing.T, mutate func(*DatabaseConfig)) *PostgresCheckpointer { + t.Helper() + + cfg := newPostgresSchema(t, mutate) + cp, err := NewPostgresCheckpointer(cfg) + require.NoError(t, err, "create PostgresCheckpointer") + t.Cleanup(func() { + if err := cp.Close(); err != nil { + t.Logf("failed to close checkpointer: %v", err) + } + }) + return cp +} + +// richState builds a state covering the value shapes that actually flow through +// a graph: strings, numbers, booleans, nils, slices and nested maps. Anything +// that silently drops data on the way to storage shows up here. +func richState() *core.BaseState { + st := core.NewBaseState() + st.Set("greeting", "hello world") + st.Set("count", 42) + st.Set("ratio", 3.5) + st.Set("enabled", true) + st.Set("missing", nil) + st.Set("tags", []interface{}{"a", "b", "c"}) + st.Set("nested", map[string]interface{}{ + "inner": "value", + "deep": map[string]interface{}{"x": 1.0}, + "list": []interface{}{1.0, 2.0}, + "quoted": `he said "hi"; DROP TABLE checkpoints;--`, + }) + st.SetMetadata("origin", "integration-test") + return st +} + +// assertRichState checks a state survived a full store/load cycle. +// +// This is the assertion the whole exercise exists for: BaseState keeps its data +// in unexported fields, so before it grew MarshalJSON/UnmarshalJSON every +// checkpoint serialized as "{}" and lost everything. A round trip that returns +// the same values is the proof that no longer happens through PostgreSQL. +func assertRichState(t *testing.T, got *core.BaseState) { + t.Helper() + require.NotNil(t, got, "state must not be nil after load") + + all := got.GetAll() + assert.Equal(t, "hello world", all["greeting"]) + // JSON has one number type, so integers come back as float64. + assert.EqualValues(t, 42, all["count"]) + assert.EqualValues(t, 3.5, all["ratio"]) + assert.Equal(t, true, all["enabled"]) + assert.Nil(t, all["missing"]) + assert.Equal(t, []interface{}{"a", "b", "c"}, all["tags"]) + + nested, ok := all["nested"].(map[string]interface{}) + require.True(t, ok, "nested value should decode as a map, got %T", all["nested"]) + assert.Equal(t, "value", nested["inner"]) + assert.Equal(t, `he said "hi"; DROP TABLE checkpoints;--`, nested["quoted"], + "quotes and SQL metacharacters must survive verbatim") + deep, ok := nested["deep"].(map[string]interface{}) + require.True(t, ok, "deeply nested map should survive") + assert.EqualValues(t, 1, deep["x"]) + + origin, found := got.GetMetadata("origin") + assert.True(t, found, "state metadata must survive the round trip") + assert.Equal(t, "integration-test", origin) +} + +func newCheckpoint(threadID, id string, step int, state *core.BaseState) *Checkpoint { + return &Checkpoint{ + ID: id, + ThreadID: threadID, + State: state, + Metadata: map[string]interface{}{"node": "n" + strconv.Itoa(step), "step": step}, + // Postgres stores microsecond precision; truncating keeps equality + // assertions honest rather than comparing against rounded values. + CreatedAt: time.Now().UTC().Truncate(time.Microsecond), + NodeID: "node-" + strconv.Itoa(step), + StepID: step, + } +} + +// --- core round trip ------------------------------------------------------ + +func TestPostgresCheckpointer_SaveLoadRoundTripPreservesState(t *testing.T) { + cp := newPostgresCheckpointer(t, nil) + ctx := context.Background() + + want := newCheckpoint("thread-round-trip", "ckpt-1", 1, richState()) + require.NoError(t, cp.Save(ctx, want)) + + got, err := cp.Load(ctx, want.ThreadID, want.ID) + require.NoError(t, err) + + assert.Equal(t, want.ID, got.ID) + assert.Equal(t, want.ThreadID, got.ThreadID) + assert.Equal(t, want.NodeID, got.NodeID) + assert.Equal(t, want.StepID, got.StepID) + assert.WithinDuration(t, want.CreatedAt, got.CreatedAt, time.Millisecond) + assert.Equal(t, "n1", got.Metadata["node"]) + assertRichState(t, got.State) +} + +// TestPostgresCheckpointer_StoresRealJSONNotEmptyObject guards the specific +// regression that motivated these tests: state used to reach the database as +// literal "{}". Reading the column back as text proves the payload really +// contains the data, independent of how Load decodes it. +func TestPostgresCheckpointer_StoresRealJSONNotEmptyObject(t *testing.T) { + cp := newPostgresCheckpointer(t, nil) + ctx := context.Background() + + ck := newCheckpoint("thread-json", "ckpt-json", 1, richState()) + require.NoError(t, cp.Save(ctx, ck)) + + row, err := asSQLRow(cp.conn.QueryRow(ctx, + `SELECT state_data::text FROM checkpoints WHERE id = $1`, ck.ID)) + require.NoError(t, err) + + var raw string + require.NoError(t, row.Scan(&raw)) + + assert.NotEqual(t, "{}", raw, "state was serialized as an empty object -- all state lost") + assert.Contains(t, raw, "hello world") + assert.Contains(t, raw, "\"count\"") + assert.Contains(t, raw, "integration-test", "state metadata must be persisted too") +} + +// TestPostgresCheckpointer_SaveCreatesThreadRow covers the defect that made the +// PostgreSQL backend unusable: checkpoints.thread_id has a FOREIGN KEY to +// threads(id), but the Checkpointer interface offers no way to create a thread, +// so every first Save failed with a foreign key violation. +func TestPostgresCheckpointer_SaveCreatesThreadRow(t *testing.T) { + cp := newPostgresCheckpointer(t, nil) + ctx := context.Background() + + ck := newCheckpoint("brand-new-thread", "ckpt-fk", 1, core.NewBaseState()) + require.NoError(t, cp.Save(ctx, ck), + "saving to a thread that was never registered must succeed, as it does for the memory and file backends") + + sm := NewSessionManager(cp.conn) + thread, err := sm.GetThread(ctx, ck.ThreadID) + require.NoError(t, err, "Save should have registered the parent thread") + assert.Equal(t, ck.ThreadID, thread.ID) +} + +func TestPostgresCheckpointer_SaveIsIdempotentUpsert(t *testing.T) { + cp := newPostgresCheckpointer(t, nil) + ctx := context.Background() + + first := core.NewBaseState() + first.Set("version", "one") + ck := newCheckpoint("thread-upsert", "ckpt-same-id", 1, first) + require.NoError(t, cp.Save(ctx, ck)) + + second := core.NewBaseState() + second.Set("version", "two") + ck.State = second + ck.NodeID = "updated" + require.NoError(t, cp.Save(ctx, ck), "re-saving the same ID must upsert, not conflict") + + got, err := cp.Load(ctx, ck.ThreadID, ck.ID) + require.NoError(t, err) + assert.Equal(t, "two", got.State.GetAll()["version"], "upsert must overwrite the stored state") + assert.Equal(t, "updated", got.NodeID) + + list, err := cp.List(ctx, ck.ThreadID) + require.NoError(t, err) + assert.Len(t, list, 1, "upsert must not create a duplicate row") +} + +// --- isolation ------------------------------------------------------------ + +func TestPostgresCheckpointer_ThreadIsolation(t *testing.T) { + cp := newPostgresCheckpointer(t, nil) + ctx := context.Background() + + a := core.NewBaseState() + a.Set("owner", "alice") + b := core.NewBaseState() + b.Set("owner", "bob") + + require.NoError(t, cp.Save(ctx, newCheckpoint("thread-a", "shared-id", 1, a))) + require.NoError(t, cp.Save(ctx, newCheckpoint("thread-b", "other-id", 1, b))) + + gotA, err := cp.Load(ctx, "thread-a", "shared-id") + require.NoError(t, err) + assert.Equal(t, "alice", gotA.State.GetAll()["owner"]) + + // A checkpoint ID that exists, but under a different thread, must not be + // readable from this thread. + _, err = cp.Load(ctx, "thread-b", "shared-id") + assert.Error(t, err, "loading another thread's checkpoint ID must fail") + + listA, err := cp.List(ctx, "thread-a") + require.NoError(t, err) + require.Len(t, listA, 1) + assert.Equal(t, "shared-id", listA[0].ID) +} + +// TestPostgresCheckpointer_IdentifiersWithSQLMetacharacters feeds hostile +// identifiers through every query. All statements use bound parameters, so +// these must be stored and matched verbatim rather than being interpreted. +func TestPostgresCheckpointer_IdentifiersWithSQLMetacharacters(t *testing.T) { + cp := newPostgresCheckpointer(t, nil) + ctx := context.Background() + + hostile := `thread'; DROP TABLE checkpoints; --` + ckID := `id" OR "1"="1` + + st := core.NewBaseState() + st.Set("safe", true) + require.NoError(t, cp.Save(ctx, newCheckpoint(hostile, ckID, 1, st))) + + got, err := cp.Load(ctx, hostile, ckID) + require.NoError(t, err, "identifiers must round trip verbatim") + assert.Equal(t, hostile, got.ThreadID) + assert.Equal(t, ckID, got.ID) + + // The table is obviously still there if this works. + list, err := cp.List(ctx, hostile) + require.NoError(t, err) + assert.Len(t, list, 1) + + // And the injected predicate must not have widened the match. + _, err = cp.Load(ctx, hostile, "no-such-checkpoint") + assert.Error(t, err) +} + +// --- listing and deletion ------------------------------------------------- + +func TestPostgresCheckpointer_ListReturnsEveryCheckpoint(t *testing.T) { + cp := newPostgresCheckpointer(t, nil) + ctx := context.Background() + + const n = 5 + for i := 0; i < n; i++ { + st := core.NewBaseState() + st.Set("i", i) + require.NoError(t, cp.Save(ctx, newCheckpoint("thread-list", fmt.Sprintf("ckpt-%d", i), i, st))) + } + + list, err := cp.List(ctx, "thread-list") + require.NoError(t, err) + require.Len(t, list, n) + + seen := map[string]bool{} + for _, m := range list { + seen[m.ID] = true + assert.Equal(t, "thread-list", m.ThreadID) + assert.NotNil(t, m.Metadata, "metadata must never be nil after List") + } + assert.Len(t, seen, n, "every checkpoint must appear exactly once") + + empty, err := cp.List(ctx, "thread-that-does-not-exist") + require.NoError(t, err, "listing an unknown thread is not an error") + assert.Empty(t, empty) +} + +func TestPostgresCheckpointer_Delete(t *testing.T) { + cp := newPostgresCheckpointer(t, nil) + ctx := context.Background() + + require.NoError(t, cp.Save(ctx, newCheckpoint("thread-del", "ckpt-del", 1, core.NewBaseState()))) + require.NoError(t, cp.Delete(ctx, "thread-del", "ckpt-del")) + + _, err := cp.Load(ctx, "thread-del", "ckpt-del") + assert.Error(t, err, "deleted checkpoint must be gone") + + list, err := cp.List(ctx, "thread-del") + require.NoError(t, err) + assert.Empty(t, list) +} + +// TestPostgresCheckpointer_DeleteMissingReportsNotFound covers a defect where +// deleting a checkpoint that did not exist reported success, so a typo'd ID +// looked like a completed deletion. The memory and file backends both error. +func TestPostgresCheckpointer_DeleteMissingReportsNotFound(t *testing.T) { + cp := newPostgresCheckpointer(t, nil) + ctx := context.Background() + + require.NoError(t, cp.Save(ctx, newCheckpoint("thread-del2", "real", 1, core.NewBaseState()))) + + err := cp.Delete(ctx, "thread-del2", "never-existed") + assert.Error(t, err, "deleting a missing checkpoint must not report success") + assert.Contains(t, err.Error(), "not found") + + // Deleting under the wrong thread must not delete the real row either. + assert.Error(t, cp.Delete(ctx, "some-other-thread", "real")) + _, err = cp.Load(ctx, "thread-del2", "real") + assert.NoError(t, err, "a mis-targeted delete must not remove the real checkpoint") +} + +func TestPostgresCheckpointer_LoadMissingIsAnError(t *testing.T) { + cp := newPostgresCheckpointer(t, nil) + ctx := context.Background() + + _, err := cp.Load(ctx, "no-thread", "no-checkpoint") + require.Error(t, err) + assert.Contains(t, err.Error(), "not found") + assert.NotErrorIs(t, err, sql.ErrNoRows, "the raw driver sentinel should be translated") +} + +// --- damaged and legacy rows ---------------------------------------------- + +// TestPostgresCheckpointer_TolerantOfNullColumns covers rows written by an older +// release or another tool. metadata, node_id and step_id are all nullable in the +// schema, and scanning a NULL used to abort Load *and* List -- and a failing +// List breaks Latest(), i.e. resuming the thread at all. +func TestPostgresCheckpointer_TolerantOfNullColumns(t *testing.T) { + cp := newPostgresCheckpointer(t, nil) + ctx := context.Background() + + ck := newCheckpoint("thread-null", "ckpt-null", 7, richState()) + require.NoError(t, cp.Save(ctx, ck)) + + require.NoError(t, cp.conn.ExecuteQuery(ctx, + `UPDATE checkpoints SET metadata = NULL, node_id = NULL, step_id = NULL WHERE id = $1`, ck.ID)) + + got, err := cp.Load(ctx, ck.ThreadID, ck.ID) + require.NoError(t, err, "a row with NULL metadata/node_id/step_id must still load") + assert.NotNil(t, got.Metadata, "NULL metadata must decode to an empty map, not nil") + assert.Empty(t, got.Metadata) + assert.Equal(t, "", got.NodeID) + assert.Equal(t, 0, got.StepID) + assertRichState(t, got.State) + + list, err := cp.List(ctx, ck.ThreadID) + require.NoError(t, err, "List must survive NULL columns as well") + require.Len(t, list, 1) + assert.NotNil(t, list[0].Metadata) +} + +// TestPostgresCheckpointer_CorruptStateIsReported checks that unreadable data is +// surfaced rather than silently returning an empty state, which would look like +// a successful resume from nothing. +func TestPostgresCheckpointer_CorruptStateIsReported(t *testing.T) { + cp := newPostgresCheckpointer(t, nil) + ctx := context.Background() + + ck := newCheckpoint("thread-corrupt", "ckpt-corrupt", 1, richState()) + require.NoError(t, cp.Save(ctx, ck)) + + // A JSON scalar where an object is expected: valid JSONB, wrong shape. + require.NoError(t, cp.conn.ExecuteQuery(ctx, + `UPDATE checkpoints SET state_data = '"not-an-object"'::jsonb WHERE id = $1`, ck.ID)) + + _, err := cp.Load(ctx, ck.ThreadID, ck.ID) + require.Error(t, err, "corrupt state must be reported, not silently dropped") + assert.Contains(t, err.Error(), "unmarshal state") +} + +// --- concurrency ---------------------------------------------------------- + +// TestPostgresCheckpointer_ConcurrentAccess drives real concurrent traffic +// through one pool. Run under -race it also covers the checkpointer's own +// shared state. +func TestPostgresCheckpointer_ConcurrentAccess(t *testing.T) { + cp := newPostgresCheckpointer(t, nil) + ctx := context.Background() + + const workers = 8 + const perWorker = 5 + + var wg sync.WaitGroup + errCh := make(chan error, workers*perWorker*2) + + for w := 0; w < workers; w++ { + wg.Add(1) + go func(w int) { + defer wg.Done() + threadID := fmt.Sprintf("concurrent-%d", w) + for i := 0; i < perWorker; i++ { + st := core.NewBaseState() + st.Set("worker", w) + st.Set("index", i) + ck := newCheckpoint(threadID, fmt.Sprintf("ckpt-%d-%d", w, i), i, st) + if err := cp.Save(ctx, ck); err != nil { + errCh <- fmt.Errorf("save: %w", err) + continue + } + got, err := cp.Load(ctx, threadID, ck.ID) + if err != nil { + errCh <- fmt.Errorf("load: %w", err) + continue + } + if got.State.GetAll()["worker"] != float64(w) { + errCh <- fmt.Errorf("worker %d read back %v", w, got.State.GetAll()["worker"]) + } + } + }(w) + } + + wg.Wait() + close(errCh) + for err := range errCh { + t.Error(err) + } + + // Every worker's thread must hold exactly its own checkpoints. + for w := 0; w < workers; w++ { + list, err := cp.List(ctx, fmt.Sprintf("concurrent-%d", w)) + require.NoError(t, err) + assert.Len(t, list, perWorker, "thread %d lost or gained checkpoints", w) + } +} + +// --- context propagation -------------------------------------------------- + +// TestPostgresCheckpointer_HonorsContextCancellation proves ctx actually +// reaches the driver rather than being accepted and ignored. +func TestPostgresCheckpointer_HonorsContextCancellation(t *testing.T) { + cp := newPostgresCheckpointer(t, nil) + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + err := cp.Save(ctx, newCheckpoint("thread-ctx", "ckpt-ctx", 1, core.NewBaseState())) + require.Error(t, err) + assert.ErrorIs(t, err, context.Canceled) + + _, err = cp.List(ctx, "thread-ctx") + require.Error(t, err) + assert.ErrorIs(t, err, context.Canceled) + + _, err = cp.Load(ctx, "thread-ctx", "ckpt-ctx") + require.Error(t, err) + assert.ErrorIs(t, err, context.Canceled) + + err = cp.Delete(ctx, "thread-ctx", "ckpt-ctx") + require.Error(t, err) + assert.ErrorIs(t, err, context.Canceled) +} + +// --- integration with the graph-facing helpers ---------------------------- + +// TestPostgresCheckpointer_SaverAndLatest exercises the path a running graph +// actually takes: CheckpointSaver writes a checkpoint per step, Latest resumes +// from the newest one. +func TestPostgresCheckpointer_SaverAndLatest(t *testing.T) { + cp := newPostgresCheckpointer(t, nil) + ctx := context.Background() + + saver := NewCheckpointSaver(cp) + threadID := "thread-saver" + + for step := 1; step <= 4; step++ { + st := core.NewBaseState() + st.Set("step", step) + st.Set("payload", fmt.Sprintf("after-node-%d", step)) + require.NoError(t, saver.SaveState(ctx, threadID, fmt.Sprintf("node-%d", step), step, st)) + } + + latest, err := Latest(ctx, cp, threadID) + require.NoError(t, err) + require.NotNil(t, latest, "a thread with checkpoints must have a latest one") + assert.EqualValues(t, 4, latest.State.GetAll()["step"], "Latest must return the highest step") + assert.Equal(t, "after-node-4", latest.State.GetAll()["payload"]) + + list, err := cp.List(ctx, threadID) + require.NoError(t, err) + assert.Len(t, list, 4, "one checkpoint per step") + + none, err := Latest(ctx, cp, "thread-with-nothing") + require.NoError(t, err) + assert.Nil(t, none) +} + +// --- RAG document storage ------------------------------------------------- + +// TestPostgresCheckpointer_DocumentRoundTrip covers the non-vector RAG path. +// SaveDocument passed its metadata map straight to database/sql, which rejects +// maps -- so this method failed 100% of the time before the fix. +func TestPostgresCheckpointer_DocumentRoundTrip(t *testing.T) { + cp := newPostgresCheckpointer(t, func(c *DatabaseConfig) { c.EnableRAG = true }) + ctx := context.Background() + + threadID := "thread-docs" + doc := &Document{ + ID: "doc-1", + ThreadID: threadID, + Content: "the quick brown fox", + Metadata: map[string]interface{}{"source": "unit-test", "page": 3.0}, + CreatedAt: time.Now().UTC().Truncate(time.Microsecond), + UpdatedAt: time.Now().UTC().Truncate(time.Microsecond), + } + require.NoError(t, cp.SaveDocument(ctx, doc), "SaveDocument must accept a metadata map") + + docs, err := cp.SearchDocuments(ctx, threadID, nil, 10) + require.NoError(t, err) + require.Len(t, docs, 1) + assert.Equal(t, "doc-1", docs[0].ID) + assert.Equal(t, "the quick brown fox", docs[0].Content) + assert.Equal(t, "unit-test", docs[0].Metadata["source"], "document metadata must round trip") + assert.EqualValues(t, 3, docs[0].Metadata["page"]) + + // Upsert on the same ID. + doc.Content = "updated content" + require.NoError(t, cp.SaveDocument(ctx, doc)) + docs, err = cp.SearchDocuments(ctx, threadID, nil, 10) + require.NoError(t, err) + require.Len(t, docs, 1, "re-saving the same document ID must upsert") + assert.Equal(t, "updated content", docs[0].Content) + + // Documents from another thread must not leak into this one. + other := *doc + other.ID = "doc-other" + other.ThreadID = "different-thread" + require.NoError(t, cp.SaveDocument(ctx, &other)) + docs, err = cp.SearchDocuments(ctx, threadID, nil, 10) + require.NoError(t, err) + assert.Len(t, docs, 1, "SearchDocuments must be scoped to its thread") +} + +func TestPostgresCheckpointer_DocumentNullMetadata(t *testing.T) { + cp := newPostgresCheckpointer(t, func(c *DatabaseConfig) { c.EnableRAG = true }) + ctx := context.Background() + + // Register the thread, then insert a row with NULL metadata the way an + // external writer would. + require.NoError(t, cp.Save(ctx, newCheckpoint("thread-doc-null", "c", 1, core.NewBaseState()))) + require.NoError(t, cp.conn.ExecuteQuery(ctx, + `INSERT INTO documents (id, thread_id, content, metadata, created_at, updated_at) + VALUES ('legacy-doc', $1, 'content', NULL, NOW(), NOW())`, "thread-doc-null")) + + docs, err := cp.SearchDocuments(ctx, "thread-doc-null", nil, 10) + require.NoError(t, err, "a document with NULL metadata must not break the search") + require.Len(t, docs, 1) + assert.NotNil(t, docs[0].Metadata) + assert.Empty(t, docs[0].Metadata) +} + +func TestPostgresCheckpointer_RAGDisabledIsRejected(t *testing.T) { + cp := newPostgresCheckpointer(t, nil) // EnableRAG defaults to false + ctx := context.Background() + + err := cp.SaveDocument(ctx, &Document{ID: "d", ThreadID: "t", Content: "c"}) + require.Error(t, err) + assert.Contains(t, err.Error(), "RAG is not enabled") + + _, err = cp.SearchDocuments(ctx, "t", nil, 5) + require.Error(t, err) + assert.Contains(t, err.Error(), "RAG is not enabled") +} + +// --- pgvector ------------------------------------------------------------- + +// TestPgVectorCheckpointer_EmbeddingRoundTrip covers the vector path against a +// real pgvector install. Two separate defects lived here: []float64 is not a +// valid driver argument (so both writes and vector searches always failed), and +// the scanned embedding was thrown away behind a "handle conversion if needed" +// comment, so every document read back had a nil Embedding. +func TestPgVectorCheckpointer_EmbeddingRoundTrip(t *testing.T) { + requirePostgres(t) + if !pgHasVector { + t.Skip("pgvector extension is not installed in the test database; " + + "run CREATE EXTENSION vector to enable this test") + } + + cp := newPostgresCheckpointer(t, func(c *DatabaseConfig) { + c.Type = DatabaseTypePgVector + c.EnableRAG = true + c.VectorDimension = 3 + }) + ctx := context.Background() + + threadID := "thread-vec" + near := &Document{ + ID: "near", ThreadID: threadID, Content: "close match", + Metadata: map[string]interface{}{"kind": "near"}, + // Deliberately non-integral so a lossy encoder would show up. + Embedding: []float64{1.0, 0.25, -0.5}, + CreatedAt: time.Now().UTC(), UpdatedAt: time.Now().UTC(), + } + far := &Document{ + ID: "far", ThreadID: threadID, Content: "distant match", + Metadata: map[string]interface{}{"kind": "far"}, + Embedding: []float64{-9, -9, -9}, + CreatedAt: time.Now().UTC(), UpdatedAt: time.Now().UTC(), + } + require.NoError(t, cp.SaveDocument(ctx, near), "SaveDocument must accept a []float64 embedding") + require.NoError(t, cp.SaveDocument(ctx, far)) + + docs, err := cp.SearchDocuments(ctx, threadID, []float64{1.0, 0.25, -0.5}, 10) + require.NoError(t, err, "vector similarity search must accept a []float64 query") + require.Len(t, docs, 2) + + // Nearest first: this is what makes it a similarity search rather than an + // arbitrary listing. + assert.Equal(t, "near", docs[0].ID, "results must be ordered by vector distance") + assert.Equal(t, "far", docs[1].ID) + + assert.Equal(t, []float64{1.0, 0.25, -0.5}, docs[0].Embedding, + "the stored embedding must be decoded back, not discarded") + assert.Equal(t, "near", docs[0].Metadata["kind"]) +} + +// --- SessionManager against a real server --------------------------------- + +func TestSessionManager_ThreadAndSessionRoundTrip(t *testing.T) { + cp := newPostgresCheckpointer(t, nil) + ctx := context.Background() + sm := NewSessionManager(cp.conn) + + thread := &Thread{ + ID: "sm-thread", + Name: "conversation", + Metadata: map[string]interface{}{"locale": "en"}, + CreatedAt: time.Now().UTC().Truncate(time.Microsecond), + UpdatedAt: time.Now().UTC().Truncate(time.Microsecond), + } + require.NoError(t, sm.CreateThread(ctx, thread)) + + gotThread, err := sm.GetThread(ctx, thread.ID) + require.NoError(t, err) + assert.Equal(t, thread.Name, gotThread.Name) + assert.Equal(t, "en", gotThread.Metadata["locale"]) + + expires := time.Now().UTC().Add(time.Hour).Truncate(time.Microsecond) + session := &Session{ + ID: "sm-session", + ThreadID: thread.ID, + UserID: "user-1", + Metadata: map[string]interface{}{"agent": "test"}, + CreatedAt: time.Now().UTC().Truncate(time.Microsecond), + ExpiresAt: &expires, + } + require.NoError(t, sm.CreateSession(ctx, session)) + + gotSession, err := sm.GetSession(ctx, session.ID) + require.NoError(t, err) + assert.Equal(t, "user-1", gotSession.UserID) + assert.Equal(t, thread.ID, gotSession.ThreadID) + assert.Equal(t, "test", gotSession.Metadata["agent"]) + require.NotNil(t, gotSession.ExpiresAt) + assert.WithinDuration(t, expires, *gotSession.ExpiresAt, time.Millisecond) + + _, err = sm.GetThread(ctx, "missing") + assert.Error(t, err) + _, err = sm.GetSession(ctx, "missing") + assert.Error(t, err) +} + +// --- connection handling -------------------------------------------------- + +// TestPostgresConnection_RejectsInvalidMaxLifetime covers a silent failure: an +// unparseable duration used to be swallowed, leaving connections with no +// lifetime cap instead of the configured one. +func TestPostgresConnection_RejectsInvalidMaxLifetime(t *testing.T) { + cfg := requirePostgres(t) + cfg.MaxLifetime = "5 minutes" // not a Go duration + + _, err := NewPostgresConnection(cfg) + require.Error(t, err, "an unparseable max_lifetime must be reported, not ignored") + assert.Contains(t, err.Error(), "max_lifetime") +} + +func TestPostgresConnection_ExecAndTransaction(t *testing.T) { + cfg := newPostgresSchema(t, nil) + conn, err := NewPostgresConnection(cfg) + require.NoError(t, err) + defer func() { require.NoError(t, conn.Close()) }() + + ctx := context.Background() + require.NoError(t, conn.ExecuteQuery(ctx, `CREATE TABLE t (id int primary key)`)) + + res, err := conn.Exec(ctx, `INSERT INTO t (id) VALUES (1), (2)`) + require.NoError(t, err) + affected, err := res.RowsAffected() + require.NoError(t, err) + assert.EqualValues(t, 2, affected) + + // A transaction that returns an error must leave nothing behind. + wantErr := fmt.Errorf("deliberate failure") + err = conn.WithTx(ctx, func(tx *sql.Tx) error { + if _, e := tx.ExecContext(ctx, `INSERT INTO t (id) VALUES (3)`); e != nil { + return e + } + return wantErr + }) + assert.ErrorIs(t, err, wantErr) + + row, err := asSQLRow(conn.QueryRow(ctx, `SELECT count(*) FROM t`)) + require.NoError(t, err) + var count int + require.NoError(t, row.Scan(&count)) + assert.Equal(t, 2, count, "a rolled back transaction must not persist its rows") +} + +func TestDatabaseConnectionManager_RealConnectionLifecycle(t *testing.T) { + cfg := newPostgresSchema(t, nil) + mgr := NewDatabaseConnectionManager() + + require.NoError(t, mgr.AddConnection("primary", cfg)) + + conn, err := mgr.GetConnection("primary") + require.NoError(t, err) + require.NoError(t, conn.Ping()) + assert.Equal(t, DatabaseTypePostgres, conn.GetType()) + + // Re-adding under the same name must close the previous pool rather than + // leaking it, and must hand back the new one. + require.NoError(t, mgr.AddConnection("primary", cfg)) + replaced, err := mgr.GetConnection("primary") + require.NoError(t, err) + assert.NotSame(t, conn, replaced, "re-adding a name must install a new connection") + assert.Error(t, conn.Ping(), "the replaced connection must have been closed") + + require.NoError(t, mgr.CloseAll()) + // CloseAll clears the registry, so a second call is a safe no-op rather + // than a double close. + require.NoError(t, mgr.CloseAll()) + _, err = mgr.GetConnection("primary") + assert.Error(t, err) +} + +// TestDatabaseConnectionManager_ConcurrentUse would crash the process with a +// "concurrent map read and map write" fatal error before the manager grew a +// mutex. Run under -race it also reports the data race itself. +func TestDatabaseConnectionManager_ConcurrentUse(t *testing.T) { + cfg := newPostgresSchema(t, nil) + mgr := NewDatabaseConnectionManager() + t.Cleanup(func() { + if err := mgr.CloseAll(); err != nil { + t.Logf("CloseAll: %v", err) + } + }) + + var wg sync.WaitGroup + for i := 0; i < 4; i++ { + wg.Add(1) + go func(i int) { + defer wg.Done() + name := fmt.Sprintf("conn-%d", i) + if err := mgr.AddConnection(name, cfg); err != nil { + t.Errorf("AddConnection: %v", err) + return + } + if _, err := mgr.GetConnection(name); err != nil { + t.Errorf("GetConnection: %v", err) + } + }(i) + } + for i := 0; i < 4; i++ { + wg.Add(1) + go func() { + defer wg.Done() + // Misses are expected; the point is that a concurrent read against + // a concurrent write must not crash. + _, _ = mgr.GetConnection("conn-0") + }() + } + wg.Wait() +} diff --git a/pkg/persistence/redis_integration_test.go b/pkg/persistence/redis_integration_test.go new file mode 100644 index 0000000..28dfcd7 --- /dev/null +++ b/pkg/persistence/redis_integration_test.go @@ -0,0 +1,599 @@ +// Copyright (c) 2024 GoLangGraph Team +// +// Licensed under the MIT License. See LICENSE file in the project root for full license information. +// +// Package: GoLangGraph - A powerful Go framework for building AI agent workflows + +// Integration tests that run RedisCheckpointer against a REAL Redis server. +// The client, the key layout, the TTLs and the JSON payloads are all exercised +// for real; nothing is stubbed. +// +// The server is taken from REDIS_TEST_ADDR when set (an unreachable explicit +// address is a hard failure), otherwise 127.0.0.1:6379 is probed and every test +// here skips cleanly if nothing answers. +// +// Redis has no schemas, so isolation comes from a per-test thread-ID prefix and +// a cleanup that deletes exactly the keys a test created. No test ever flushes +// the database, which would destroy unrelated data on a shared server. + +package persistence + +import ( + "context" + "fmt" + "os" + "strings" + "sync" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/UnicoLab/GoLangGraph/pkg/core" +) + +const defaultRedisTestAddr = "127.0.0.1:6379" + +var ( + redisDiscoverOnce sync.Once + redisBaseConfig *DatabaseConfig + redisDiscoverErr error + redisExplicitAddr bool +) + +func parseRedisAddr(addr string) (*DatabaseConfig, error) { + host, portStr, found := strings.Cut(addr, ":") + if !found { + return nil, fmt.Errorf("invalid redis address %q, want host:port", addr) + } + var port int + if _, err := fmt.Sscanf(portStr, "%d", &port); err != nil { + return nil, fmt.Errorf("invalid port in redis address %q: %w", addr, err) + } + // Password comes from REDIS_TEST_PASSWORD; sending a password to a server + // with none configured is itself an error, so the default is empty. + return NewRedisConfig(host, port, os.Getenv("REDIS_TEST_PASSWORD")), nil +} + +func discoverRedis() (*DatabaseConfig, error) { + redisDiscoverOnce.Do(func() { + addr := defaultRedisTestAddr + if explicit := os.Getenv("REDIS_TEST_ADDR"); explicit != "" { + addr = explicit + redisExplicitAddr = true + } + + cfg, err := parseRedisAddr(addr) + if err != nil { + redisDiscoverErr = err + return + } + cp, err := NewRedisCheckpointer(cfg) + if err != nil { + redisDiscoverErr = err + return + } + _ = cp.Close() + redisBaseConfig = cfg + }) + return redisBaseConfig, redisDiscoverErr +} + +func requireRedisConfig(t *testing.T) *DatabaseConfig { + t.Helper() + + cfg, err := discoverRedis() + if cfg == nil { + if redisExplicitAddr { + t.Fatalf("REDIS_TEST_ADDR is set but the server is unreachable: %v", err) + } + t.Skipf("no local Redis reachable at %s (%v). "+ + "Set REDIS_TEST_ADDR to run these integration tests.", defaultRedisTestAddr, err) + } + clone := *cfg + return &clone +} + +// newRedisCheckpointer returns a checkpointer plus a prefix unique to this test. +// Cleanup removes every key under that prefix, so a shared Redis is left exactly +// as it was found. +func newRedisCheckpointer(t *testing.T, mutate func(*DatabaseConfig)) (*RedisCheckpointer, string) { + t.Helper() + + cfg := requireRedisConfig(t) + if mutate != nil { + mutate(cfg) + } + + cp, err := NewRedisCheckpointer(cfg) + require.NoError(t, err, "connect to Redis") + + prefix := fmt.Sprintf("gltest:%s:%d:", t.Name(), time.Now().UnixNano()) + + t.Cleanup(func() { + ctx := context.Background() + // Match both the checkpoint payloads and the thread index sets. The + // prefix contains ':' so it is escaped in keys exactly as the + // checkpointer escapes it. + for _, pattern := range []string{ + "checkpoint:" + redisKeySegment(prefix) + "*", + "thread:" + redisKeySegment(prefix) + "*", + } { + keys, err := cp.client.Keys(ctx, pattern).Result() + if err != nil { + t.Logf("cleanup scan %q: %v", pattern, err) + continue + } + if len(keys) > 0 { + if err := cp.client.Del(ctx, keys...).Err(); err != nil { + t.Logf("cleanup delete: %v", err) + } + } + } + if err := cp.Close(); err != nil { + t.Logf("failed to close redis checkpointer: %v", err) + } + }) + + return cp, prefix +} + +// --- core round trip ------------------------------------------------------ + +// TestRedisCheckpointer_SaveLoadRoundTripPreservesState is the reason this file +// exists. BaseState stores its data in unexported fields, so before it gained +// MarshalJSON/UnmarshalJSON every checkpoint went to Redis as "{}" and came +// back empty. This proves a real server returns the real state. +func TestRedisCheckpointer_SaveLoadRoundTripPreservesState(t *testing.T) { + cp, prefix := newRedisCheckpointer(t, nil) + ctx := context.Background() + + want := newCheckpoint(prefix+"round-trip", "ckpt-1", 3, richState()) + require.NoError(t, cp.Save(ctx, want)) + + got, err := cp.Load(ctx, want.ThreadID, want.ID) + require.NoError(t, err) + + assert.Equal(t, want.ID, got.ID) + assert.Equal(t, want.ThreadID, got.ThreadID) + assert.Equal(t, want.NodeID, got.NodeID) + assert.Equal(t, want.StepID, got.StepID) + assert.WithinDuration(t, want.CreatedAt, got.CreatedAt, time.Millisecond) + assert.Equal(t, "n3", got.Metadata["node"]) + assertRichState(t, got.State) +} + +// TestRedisCheckpointer_StoresRealJSONNotEmptyObject inspects the stored bytes +// directly, so it holds even if Load were changed to fabricate a state. +func TestRedisCheckpointer_StoresRealJSONNotEmptyObject(t *testing.T) { + cp, prefix := newRedisCheckpointer(t, nil) + ctx := context.Background() + + ck := newCheckpoint(prefix+"json", "ckpt-json", 1, richState()) + require.NoError(t, cp.Save(ctx, ck)) + + raw, err := cp.client.Get(ctx, redisCheckpointKey(ck.ThreadID, ck.ID)).Result() + require.NoError(t, err) + + assert.NotContains(t, raw, `"state":{}`, "state was stored as an empty object -- all state lost") + assert.Contains(t, raw, "hello world") + assert.Contains(t, raw, "integration-test", "state metadata must be persisted too") +} + +func TestRedisCheckpointer_SaveOverwritesSameID(t *testing.T) { + cp, prefix := newRedisCheckpointer(t, nil) + ctx := context.Background() + + threadID := prefix + "overwrite" + first := core.NewBaseState() + first.Set("version", "one") + ck := newCheckpoint(threadID, "same", 1, first) + require.NoError(t, cp.Save(ctx, ck)) + + second := core.NewBaseState() + second.Set("version", "two") + ck.State = second + require.NoError(t, cp.Save(ctx, ck)) + + got, err := cp.Load(ctx, threadID, "same") + require.NoError(t, err) + assert.Equal(t, "two", got.State.GetAll()["version"]) + + list, err := cp.List(ctx, threadID) + require.NoError(t, err) + assert.Len(t, list, 1, "the thread index must not gain a duplicate entry") +} + +// --- isolation ------------------------------------------------------------ + +func TestRedisCheckpointer_ThreadIsolation(t *testing.T) { + cp, prefix := newRedisCheckpointer(t, nil) + ctx := context.Background() + + a := core.NewBaseState() + a.Set("owner", "alice") + b := core.NewBaseState() + b.Set("owner", "bob") + + require.NoError(t, cp.Save(ctx, newCheckpoint(prefix+"thread-a", "shared-id", 1, a))) + require.NoError(t, cp.Save(ctx, newCheckpoint(prefix+"thread-b", "shared-id", 1, b))) + + gotA, err := cp.Load(ctx, prefix+"thread-a", "shared-id") + require.NoError(t, err) + assert.Equal(t, "alice", gotA.State.GetAll()["owner"], "threads sharing a checkpoint ID must not share data") + + gotB, err := cp.Load(ctx, prefix+"thread-b", "shared-id") + require.NoError(t, err) + assert.Equal(t, "bob", gotB.State.GetAll()["owner"]) + + listA, err := cp.List(ctx, prefix+"thread-a") + require.NoError(t, err) + assert.Len(t, listA, 1, "one thread's index must not include the other's checkpoints") +} + +// TestRedisCheckpointer_ColonInIdentifiersDoesNotCollide covers a cross-thread +// data leak. Keys were built as fmt.Sprintf("checkpoint:%s:%s", threadID, id), +// so thread "x:a" + checkpoint "b:c1" and thread "x:a:b" + checkpoint "c1" both +// produced the key "checkpoint:x:a:b:c1": each thread read and overwrote the +// other's state. Thread IDs are routinely derived from user or session +// identifiers, which makes this a tenant-isolation failure, not a curiosity. +func TestRedisCheckpointer_ColonInIdentifiersDoesNotCollide(t *testing.T) { + cp, prefix := newRedisCheckpointer(t, nil) + ctx := context.Background() + + stateA := core.NewBaseState() + stateA.Set("which", "A") + stateB := core.NewBaseState() + stateB.Set("which", "B") + + // These two (threadID, checkpointID) pairs concatenate to the same string. + a := newCheckpoint(prefix+"x:a", "b:c1", 1, stateA) + b := newCheckpoint(prefix+"x:a:b", "c1", 1, stateB) + + require.NoError(t, cp.Save(ctx, a)) + require.NoError(t, cp.Save(ctx, b)) + + gotA, err := cp.Load(ctx, a.ThreadID, a.ID) + require.NoError(t, err) + assert.Equal(t, "A", gotA.State.GetAll()["which"], + "checkpoint A was overwritten by B -- colliding Redis keys leak state across threads") + + gotB, err := cp.Load(ctx, b.ThreadID, b.ID) + require.NoError(t, err) + assert.Equal(t, "B", gotB.State.GetAll()["which"]) + + // The two thread indexes must also stay separate. + listA, err := cp.List(ctx, a.ThreadID) + require.NoError(t, err) + require.Len(t, listA, 1) + assert.Equal(t, "b:c1", listA[0].ID) + + listB, err := cp.List(ctx, b.ThreadID) + require.NoError(t, err) + require.Len(t, listB, 1) + assert.Equal(t, "c1", listB[0].ID) +} + +// TestRedisKeySegment_EscapingIsUnambiguous checks the escaping directly, and +// pins the property that ordinary identifiers keep their original key bytes so +// data written before the fix is still readable. +func TestRedisKeySegment_EscapingIsUnambiguous(t *testing.T) { + assert.Equal(t, "plain-id_123", redisKeySegment("plain-id_123"), + "identifiers without ':' or '%' must be left byte-identical") + assert.Equal(t, "checkpoint:thread:ckpt", redisCheckpointKey("thread", "ckpt")) + + assert.NotEqual(t, + redisCheckpointKey("x:a", "b:c"), + redisCheckpointKey("x:a:b", "c"), + "identifier boundaries must survive escaping") + + // '%' is escaped too, so an attacker cannot hand-craft an identifier whose + // escaped form equals another pair's. + assert.NotEqual(t, + redisCheckpointKey("a%3Ab", "c"), + redisCheckpointKey("a:b", "c")) +} + +// --- listing and deletion ------------------------------------------------- + +func TestRedisCheckpointer_ListReturnsEveryCheckpoint(t *testing.T) { + cp, prefix := newRedisCheckpointer(t, nil) + ctx := context.Background() + + threadID := prefix + "list" + const n = 5 + for i := 0; i < n; i++ { + st := core.NewBaseState() + st.Set("i", i) + require.NoError(t, cp.Save(ctx, newCheckpoint(threadID, fmt.Sprintf("ckpt-%d", i), i, st))) + } + + list, err := cp.List(ctx, threadID) + require.NoError(t, err) + require.Len(t, list, n) + + seen := map[string]bool{} + for _, m := range list { + seen[m.ID] = true + assert.Equal(t, threadID, m.ThreadID) + } + assert.Len(t, seen, n, "every checkpoint must appear exactly once") + + empty, err := cp.List(ctx, prefix+"never-used") + require.NoError(t, err, "listing an unknown thread is not an error") + assert.Empty(t, empty) +} + +func TestRedisCheckpointer_Delete(t *testing.T) { + cp, prefix := newRedisCheckpointer(t, nil) + ctx := context.Background() + + threadID := prefix + "delete" + require.NoError(t, cp.Save(ctx, newCheckpoint(threadID, "ckpt", 1, core.NewBaseState()))) + require.NoError(t, cp.Delete(ctx, threadID, "ckpt")) + + _, err := cp.Load(ctx, threadID, "ckpt") + assert.Error(t, err) + + list, err := cp.List(ctx, threadID) + require.NoError(t, err) + assert.Empty(t, list, "Delete must also remove the thread index entry") +} + +// TestRedisCheckpointer_DeleteMissingReportsNotFound covers a defect where +// deleting something that was not there reported success. The memory and file +// backends both return an error, so callers could not rely on the result. +func TestRedisCheckpointer_DeleteMissingReportsNotFound(t *testing.T) { + cp, prefix := newRedisCheckpointer(t, nil) + ctx := context.Background() + + err := cp.Delete(ctx, prefix+"nothing-here", "ckpt") + require.Error(t, err, "deleting a missing checkpoint must not report success") + assert.Contains(t, err.Error(), "not found") +} + +// --- damaged data --------------------------------------------------------- + +// TestRedisCheckpointer_CorruptPayload pins deliberate behavior: a single +// unreadable entry must not make the whole thread unlistable (and therefore +// unresumable), but reading it directly must still report the problem. +func TestRedisCheckpointer_CorruptPayload(t *testing.T) { + cp, prefix := newRedisCheckpointer(t, nil) + ctx := context.Background() + + threadID := prefix + "corrupt" + require.NoError(t, cp.Save(ctx, newCheckpoint(threadID, "good", 1, richState()))) + + // Write a broken payload and index it the way Save would. + require.NoError(t, cp.client.Set(ctx, + redisCheckpointKey(threadID, "bad"), "{not json", time.Minute).Err()) + require.NoError(t, cp.client.SAdd(ctx, redisThreadIndexKey(threadID), "bad").Err()) + + _, err := cp.Load(ctx, threadID, "bad") + require.Error(t, err, "a corrupt payload must be reported when loaded directly") + assert.Contains(t, err.Error(), "unmarshal") + + list, err := cp.List(ctx, threadID) + require.NoError(t, err, "one bad entry must not make the thread unlistable") + require.Len(t, list, 1, "the readable checkpoint must still be returned") + assert.Equal(t, "good", list[0].ID) +} + +// TestRedisCheckpointer_StaleIndexEntryIsSkipped simulates the normal +// consequence of TTL expiry: the index still names a checkpoint whose payload +// is gone. +func TestRedisCheckpointer_StaleIndexEntryIsSkipped(t *testing.T) { + cp, prefix := newRedisCheckpointer(t, nil) + ctx := context.Background() + + threadID := prefix + "stale" + require.NoError(t, cp.Save(ctx, newCheckpoint(threadID, "alive", 1, core.NewBaseState()))) + require.NoError(t, cp.client.SAdd(ctx, redisThreadIndexKey(threadID), "expired").Err()) + + list, err := cp.List(ctx, threadID) + require.NoError(t, err) + require.Len(t, list, 1) + assert.Equal(t, "alive", list[0].ID) + + // Deleting a stale entry must clean the index even though the payload is + // already gone, so the index cannot accumulate dead members forever. + assert.Error(t, cp.Delete(ctx, threadID, "expired"), "the payload really is missing") + members, err := cp.client.SMembers(ctx, redisThreadIndexKey(threadID)).Result() + require.NoError(t, err) + assert.NotContains(t, members, "expired", "Delete must drop the stale index entry") +} + +// --- TTL handling --------------------------------------------------------- + +// TestRedisCheckpointer_ThreadIndexExpiresWithCheckpoints covers an unbounded +// leak: checkpoint payloads were written with a TTL but the thread index set was +// created with none, so it survived forever, growing a dead member per expired +// checkpoint and costing List a wasted round trip for each. +func TestRedisCheckpointer_ThreadIndexExpiresWithCheckpoints(t *testing.T) { + cp, prefix := newRedisCheckpointer(t, func(c *DatabaseConfig) { c.CheckpointTTL = "90s" }) + ctx := context.Background() + + threadID := prefix + "ttl" + require.NoError(t, cp.Save(ctx, newCheckpoint(threadID, "ckpt", 1, core.NewBaseState()))) + + payloadTTL, err := cp.client.TTL(ctx, redisCheckpointKey(threadID, "ckpt")).Result() + require.NoError(t, err) + assert.Greater(t, payloadTTL, time.Duration(0), "the checkpoint must carry the configured TTL") + assert.LessOrEqual(t, payloadTTL, 90*time.Second) + + indexTTL, err := cp.client.TTL(ctx, redisThreadIndexKey(threadID)).Result() + require.NoError(t, err) + assert.Greater(t, indexTTL, time.Duration(0), + "the thread index must expire too, or it leaks dead entries forever") + assert.LessOrEqual(t, indexTTL, 90*time.Second) +} + +// TestRedisCheckpointer_TTLIsConfigurable covers a value that was hard-coded at +// 24h with no way to change it, so every deployment silently lost its +// checkpoints after a day. +func TestRedisCheckpointer_TTLIsConfigurable(t *testing.T) { + requireRedisConfig(t) + + t.Run("default", func(t *testing.T) { + cp, _ := newRedisCheckpointer(t, nil) + assert.Equal(t, 24*time.Hour, cp.ttl) + }) + + t.Run("configured", func(t *testing.T) { + cp, prefix := newRedisCheckpointer(t, func(c *DatabaseConfig) { c.CheckpointTTL = "168h" }) + assert.Equal(t, 168*time.Hour, cp.ttl) + + ctx := context.Background() + threadID := prefix + "week" + require.NoError(t, cp.Save(ctx, newCheckpoint(threadID, "c", 1, core.NewBaseState()))) + ttl, err := cp.client.TTL(ctx, redisCheckpointKey(threadID, "c")).Result() + require.NoError(t, err) + assert.Greater(t, ttl, 167*time.Hour, "the configured TTL must reach Redis") + }) + + t.Run("zero disables expiry", func(t *testing.T) { + cp, prefix := newRedisCheckpointer(t, func(c *DatabaseConfig) { c.CheckpointTTL = "0" }) + assert.Equal(t, time.Duration(0), cp.ttl) + + ctx := context.Background() + threadID := prefix + "forever" + require.NoError(t, cp.Save(ctx, newCheckpoint(threadID, "c", 1, core.NewBaseState()))) + ttl, err := cp.client.TTL(ctx, redisCheckpointKey(threadID, "c")).Result() + require.NoError(t, err) + // Redis reports -1 for a key with no expiry, which go-redis maps to -1ns. + assert.Equal(t, time.Duration(-1), ttl, "TTL 0 must mean no expiry at all") + }) + + t.Run("invalid is rejected", func(t *testing.T) { + cfg := requireRedisConfig(t) + cfg.CheckpointTTL = "one week" + _, err := NewRedisCheckpointer(cfg) + require.Error(t, err, "an unparseable TTL must be reported, not silently ignored") + assert.Contains(t, err.Error(), "checkpoint_ttl") + }) +} + +// --- concurrency ---------------------------------------------------------- + +func TestRedisCheckpointer_ConcurrentAccess(t *testing.T) { + cp, prefix := newRedisCheckpointer(t, nil) + ctx := context.Background() + + const workers = 8 + const perWorker = 5 + + var wg sync.WaitGroup + errCh := make(chan error, workers*perWorker*2) + + for w := 0; w < workers; w++ { + wg.Add(1) + go func(w int) { + defer wg.Done() + threadID := fmt.Sprintf("%sconcurrent-%d", prefix, w) + for i := 0; i < perWorker; i++ { + st := core.NewBaseState() + st.Set("worker", w) + ck := newCheckpoint(threadID, fmt.Sprintf("ckpt-%d-%d", w, i), i, st) + if err := cp.Save(ctx, ck); err != nil { + errCh <- fmt.Errorf("save: %w", err) + continue + } + got, err := cp.Load(ctx, threadID, ck.ID) + if err != nil { + errCh <- fmt.Errorf("load: %w", err) + continue + } + if got.State.GetAll()["worker"] != float64(w) { + errCh <- fmt.Errorf("worker %d read back %v", w, got.State.GetAll()["worker"]) + } + } + }(w) + } + + wg.Wait() + close(errCh) + for err := range errCh { + t.Error(err) + } + + for w := 0; w < workers; w++ { + list, err := cp.List(ctx, fmt.Sprintf("%sconcurrent-%d", prefix, w)) + require.NoError(t, err) + assert.Len(t, list, perWorker, "thread %d lost or gained checkpoints", w) + } +} + +// --- context propagation -------------------------------------------------- + +func TestRedisCheckpointer_HonorsContextCancellation(t *testing.T) { + cp, prefix := newRedisCheckpointer(t, nil) + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + threadID := prefix + "ctx" + err := cp.Save(ctx, newCheckpoint(threadID, "c", 1, core.NewBaseState())) + require.Error(t, err) + assert.ErrorIs(t, err, context.Canceled) + + _, err = cp.Load(ctx, threadID, "c") + require.Error(t, err) + assert.ErrorIs(t, err, context.Canceled) + + _, err = cp.List(ctx, threadID) + require.Error(t, err) + assert.ErrorIs(t, err, context.Canceled) + + err = cp.Delete(ctx, threadID, "c") + require.Error(t, err) + assert.ErrorIs(t, err, context.Canceled) +} + +// --- integration with the graph-facing helpers ---------------------------- + +func TestRedisCheckpointer_SaverAndLatest(t *testing.T) { + cp, prefix := newRedisCheckpointer(t, nil) + ctx := context.Background() + + saver := NewCheckpointSaver(cp) + threadID := prefix + "saver" + + for step := 1; step <= 4; step++ { + st := core.NewBaseState() + st.Set("step", step) + st.Set("payload", fmt.Sprintf("after-node-%d", step)) + require.NoError(t, saver.SaveState(ctx, threadID, fmt.Sprintf("node-%d", step), step, st)) + } + + // List is backed by an unordered Redis set, so this also confirms Latest + // does its own ordering rather than trusting the listing order. + latest, err := Latest(ctx, cp, threadID) + require.NoError(t, err) + require.NotNil(t, latest) + assert.EqualValues(t, 4, latest.State.GetAll()["step"], "Latest must return the highest step") + assert.Equal(t, "after-node-4", latest.State.GetAll()["payload"]) + + none, err := Latest(ctx, cp, prefix+"empty") + require.NoError(t, err) + assert.Nil(t, none) +} + +// --- construction --------------------------------------------------------- + +func TestNewRedisCheckpointer_UnreachableServerFails(t *testing.T) { + // Port 1 is reserved and never listening, so this is deterministic and + // needs no server at all -- it runs even where Redis is absent. + cfg := NewRedisConfig("127.0.0.1", 1, "") + cp, err := NewRedisCheckpointer(cfg) + require.Error(t, err, "an unreachable server must fail construction") + assert.Nil(t, cp) + assert.Contains(t, err.Error(), "failed to connect to Redis") +} + +func TestRedisCheckpointer_SaveNilCheckpointIsRejected(t *testing.T) { + cp, _ := newRedisCheckpointer(t, nil) + err := cp.Save(context.Background(), nil) + require.Error(t, err, "a nil checkpoint must be rejected, not panic") + assert.Contains(t, err.Error(), "nil checkpoint") +} diff --git a/pkg/persistence/saver.go b/pkg/persistence/saver.go new file mode 100644 index 0000000..a32d2d4 --- /dev/null +++ b/pkg/persistence/saver.go @@ -0,0 +1,95 @@ +// Copyright (c) 2024 GoLangGraph Team +// +// Licensed under the MIT License. See LICENSE file in the project root for full license information. +// +// Package: GoLangGraph - A powerful Go framework for building AI agent workflows + +package persistence + +import ( + "context" + "fmt" + "sync/atomic" + "time" + + "github.com/UnicoLab/GoLangGraph/pkg/core" + "github.com/google/uuid" +) + +// CheckpointSaver adapts a Checkpointer to core.StateSaver so a graph can +// persist its state after every node execution. This is what makes execution +// durable: a run that dies mid-graph can be resumed from its last checkpoint. +// +// core deliberately does not import this package; the adapter lives here so the +// engine stays free of storage dependencies. +type CheckpointSaver struct { + checkpointer Checkpointer + seq atomic.Uint64 + // IDFunc generates checkpoint IDs. Overridable for deterministic tests. + IDFunc func(threadID, nodeID string, step int) string + // Now supplies timestamps. Overridable for deterministic tests. + Now func() time.Time +} + +// NewCheckpointSaver wraps a Checkpointer for use with core.Graph. +func NewCheckpointSaver(cp Checkpointer) *CheckpointSaver { + return &CheckpointSaver{checkpointer: cp} +} + +// SaveState implements core.StateSaver. +func (s *CheckpointSaver) SaveState(ctx context.Context, threadID, nodeID string, step int, state *core.BaseState) error { + if s == nil || s.checkpointer == nil { + return nil + } + if state == nil { + return fmt.Errorf("cannot checkpoint a nil state") + } + + now := time.Now + if s.Now != nil { + now = s.Now + } + + id := s.nextID(threadID, nodeID, step) + + return s.checkpointer.Save(ctx, &Checkpoint{ + ID: id, + ThreadID: threadID, + State: state.Clone(), + NodeID: nodeID, + StepID: step, + CreatedAt: now(), + Metadata: map[string]interface{}{ + "node_id": nodeID, + "step": step, + }, + }) +} + +func (s *CheckpointSaver) nextID(threadID, nodeID string, step int) string { + if s.IDFunc != nil { + return s.IDFunc(threadID, nodeID, step) + } + // Monotonic prefix keeps checkpoint IDs sortable by creation order, and the + // UUID suffix keeps them unique across processes sharing a thread. + return fmt.Sprintf("%010d-%s", s.seq.Add(1), uuid.New().String()) +} + +// Latest returns the most recent checkpoint for a thread, or nil when the +// thread has none. Checkpoints are ordered by step then creation time. +func Latest(ctx context.Context, cp Checkpointer, threadID string) (*Checkpoint, error) { + metas, err := cp.List(ctx, threadID) + if err != nil { + return nil, err + } + if len(metas) == 0 { + return nil, nil + } + best := metas[0] + for _, m := range metas[1:] { + if m.StepID > best.StepID || (m.StepID == best.StepID && m.CreatedAt.After(best.CreatedAt)) { + best = m + } + } + return cp.Load(ctx, threadID, best.ID) +} diff --git a/pkg/server/auto_handlers.go b/pkg/server/auto_handlers.go index cf6642e..aa5c705 100644 --- a/pkg/server/auto_handlers.go +++ b/pkg/server/auto_handlers.go @@ -29,7 +29,7 @@ func (as *AutoServer) handleHealth(w http.ResponseWriter, r *http.Request) { } w.Header().Set("Content-Type", "application/json") - json.NewEncoder(w).Encode(health) + _ = json.NewEncoder(w).Encode(health) } // handleCapabilities handles system capabilities requests @@ -56,7 +56,7 @@ func (as *AutoServer) handleCapabilities(w http.ResponseWriter, r *http.Request) } w.Header().Set("Content-Type", "application/json") - json.NewEncoder(w).Encode(capabilities) + _ = json.NewEncoder(w).Encode(capabilities) } // handleListAgents handles agent listing requests @@ -84,7 +84,7 @@ func (as *AutoServer) handleListAgents(w http.ResponseWriter, r *http.Request) { } w.Header().Set("Content-Type", "application/json") - json.NewEncoder(w).Encode(response) + _ = json.NewEncoder(w).Encode(response) } // handleAgentInfo handles individual agent information requests @@ -92,14 +92,18 @@ func (as *AutoServer) handleAgentInfo(w http.ResponseWriter, r *http.Request) { vars := mux.Vars(r) agentID := vars["agentId"] - metadata, exists := as.agentMetadata[agentID] + metadata, exists := as.agentMeta(agentID) if !exists { http.Error(w, "Agent not found", http.StatusNotFound) return } - agent := as.agentInstances[agentID] - config := agent.GetConfig() + instance, exists := as.agentInstance(agentID) + if !exists { + http.Error(w, "Agent not found", http.StatusNotFound) + return + } + config := instance.GetConfig() // Create description from system prompt or fallback description := config.SystemPrompt @@ -130,7 +134,7 @@ func (as *AutoServer) handleAgentInfo(w http.ResponseWriter, r *http.Request) { } w.Header().Set("Content-Type", "application/json") - json.NewEncoder(w).Encode(info) + _ = json.NewEncoder(w).Encode(info) } // createAgentHandler creates a handler for agent execution @@ -142,7 +146,7 @@ func (as *AutoServer) createAgentHandler(agentID string) http.HandlerFunc { return } - agent, exists := as.agentInstances[agentID] + agent, exists := as.agentInstance(agentID) if !exists { http.Error(w, "Agent not found", http.StatusNotFound) return @@ -185,7 +189,7 @@ func (as *AutoServer) createAgentHandler(agentID string) http.HandlerFunc { } w.Header().Set("Content-Type", "application/json") w.WriteHeader(http.StatusInternalServerError) - json.NewEncoder(w).Encode(response) + _ = json.NewEncoder(w).Encode(response) return } @@ -203,11 +207,24 @@ func (as *AutoServer) createAgentHandler(agentID string) http.HandlerFunc { output = result.Output } + // Report whether the output matches the schema this agent advertises, + // rather than asserting that it does. + schemaValid := true + if outputMap, ok := output.(map[string]interface{}); ok { + schema := as.generateAgentSchema(agentID, as.agentMetaOrNil(agentID)) + if issues := validateAgainstSchema(schemaSection(schema, "output"), outputMap); len(issues) > 0 { + schemaValid = false + as.logger.WithField("agent_id", agentID). + WithField("issues", summariseErrors(issues)). + Debug("Agent output did not match its advertised schema") + } + } + response := map[string]interface{}{ "success": true, "agent_id": agentID, "output": output, - "schema_valid": true, // TODO: Implement schema validation + "schema_valid": schemaValid, "processing_time": time.Since(start).String(), "timestamp": time.Now().UTC().Format(time.RFC3339), "execution_id": result.ID, @@ -216,12 +233,12 @@ func (as *AutoServer) createAgentHandler(agentID string) http.HandlerFunc { } // Add agent metadata for better frontend integration - if metadata, exists := as.agentMetadata[agentID]; exists { + if metadata, exists := as.agentMeta(agentID); exists { response["agent_metadata"] = metadata } // Add graph information if available - if agent, exists := as.agentInstances[agentID]; exists && agent.GetGraph() != nil { + if agent, exists := as.agentInstance(agentID); exists && agent.GetGraph() != nil { graph := agent.GetGraph() response["graph_info"] = map[string]interface{}{ "graph_id": graph.ID, @@ -234,7 +251,7 @@ func (as *AutoServer) createAgentHandler(agentID string) http.HandlerFunc { } w.Header().Set("Content-Type", "application/json") - json.NewEncoder(w).Encode(response) + _ = json.NewEncoder(w).Encode(response) } } @@ -247,7 +264,7 @@ func (as *AutoServer) createAgentStreamHandler(agentID string) http.HandlerFunc return } - agent, exists := as.agentInstances[agentID] + agent, exists := as.agentInstance(agentID) if !exists { http.Error(w, "Agent not found", http.StatusNotFound) return @@ -257,7 +274,11 @@ func (as *AutoServer) createAgentStreamHandler(agentID string) http.HandlerFunc w.Header().Set("Content-Type", "text/event-stream") w.Header().Set("Cache-Control", "no-cache") w.Header().Set("Connection", "keep-alive") - w.Header().Set("Access-Control-Allow-Origin", "*") + // Honor the configured allowlist rather than opening the stream to + // any origin, which the "*" here previously did. + if allow := as.allowedOrigin(r); allow != "" { + w.Header().Set("Access-Control-Allow-Origin", allow) + } // Parse request body var requestData map[string]interface{} @@ -293,7 +314,7 @@ func (as *AutoServer) createAgentStreamHandler(agentID string) http.HandlerFunc result, err := agent.Execute(ctx, input) if err != nil { - fmt.Fprintf(w, "data: {\"error\": \"%s\"}\n\n", err.Error()) + _, _ = fmt.Fprintf(w, "data: {\"error\": \"%s\"}\n\n", err.Error()) flusher.Flush() return } @@ -306,7 +327,7 @@ func (as *AutoServer) createAgentStreamHandler(agentID string) http.HandlerFunc "complete": true, }) - fmt.Fprintf(w, "data: %s\n\n", responseData) + _, _ = fmt.Fprintf(w, "data: %s\n\n", responseData) flusher.Flush() } } @@ -314,7 +335,7 @@ func (as *AutoServer) createAgentStreamHandler(agentID string) http.HandlerFunc // createConversationHandler creates a handler for conversation management func (as *AutoServer) createConversationHandler(agentID string) http.HandlerFunc { return func(w http.ResponseWriter, r *http.Request) { - agent, exists := as.agentInstances[agentID] + agent, exists := as.agentInstance(agentID) if !exists { http.Error(w, "Agent not found", http.StatusNotFound) return @@ -331,7 +352,7 @@ func (as *AutoServer) createConversationHandler(agentID string) http.HandlerFunc "timestamp": time.Now().UTC().Format(time.RFC3339), } w.Header().Set("Content-Type", "application/json") - json.NewEncoder(w).Encode(response) + _ = json.NewEncoder(w).Encode(response) case "POST": // Add to conversation @@ -349,7 +370,7 @@ func (as *AutoServer) createConversationHandler(agentID string) http.HandlerFunc "timestamp": time.Now().UTC().Format(time.RFC3339), } w.Header().Set("Content-Type", "application/json") - json.NewEncoder(w).Encode(response) + _ = json.NewEncoder(w).Encode(response) case "DELETE": // Clear conversation @@ -361,7 +382,7 @@ func (as *AutoServer) createConversationHandler(agentID string) http.HandlerFunc "timestamp": time.Now().UTC().Format(time.RFC3339), } w.Header().Set("Content-Type", "application/json") - json.NewEncoder(w).Encode(response) + _ = json.NewEncoder(w).Encode(response) default: http.Error(w, "Method not allowed", http.StatusMethodNotAllowed) @@ -372,7 +393,7 @@ func (as *AutoServer) createConversationHandler(agentID string) http.HandlerFunc // createStatusHandler creates a handler for agent status func (as *AutoServer) createStatusHandler(agentID string) http.HandlerFunc { return func(w http.ResponseWriter, r *http.Request) { - agent, exists := as.agentInstances[agentID] + agent, exists := as.agentInstance(agentID) if !exists { http.Error(w, "Agent not found", http.StatusNotFound) return @@ -388,7 +409,7 @@ func (as *AutoServer) createStatusHandler(agentID string) http.HandlerFunc { } w.Header().Set("Content-Type", "application/json") - json.NewEncoder(w).Encode(status) + _ = json.NewEncoder(w).Encode(status) } } @@ -444,7 +465,7 @@ func (as *AutoServer) handleSchemas(w http.ResponseWriter, r *http.Request) { } w.Header().Set("Content-Type", "application/json") - json.NewEncoder(w).Encode(response) + _ = json.NewEncoder(w).Encode(response) } // handleAgentSchema handles individual agent schema requests @@ -452,7 +473,7 @@ func (as *AutoServer) handleAgentSchema(w http.ResponseWriter, r *http.Request) vars := mux.Vars(r) agentID := vars["agentId"] - metadata, exists := as.agentMetadata[agentID] + metadata, exists := as.agentMeta(agentID) if !exists { http.Error(w, "Agent not found", http.StatusNotFound) return @@ -466,7 +487,7 @@ func (as *AutoServer) handleAgentSchema(w http.ResponseWriter, r *http.Request) } w.Header().Set("Content-Type", "application/json") - json.NewEncoder(w).Encode(response) + _ = json.NewEncoder(w).Encode(response) } // handleValidateSchema handles schema validation requests @@ -474,7 +495,7 @@ func (as *AutoServer) handleValidateSchema(w http.ResponseWriter, r *http.Reques vars := mux.Vars(r) agentID := vars["agentId"] - if _, exists := as.agentMetadata[agentID]; !exists { + if _, exists := as.agentMeta(agentID); !exists { http.Error(w, "Agent not found", http.StatusNotFound) return } @@ -485,22 +506,34 @@ func (as *AutoServer) handleValidateSchema(w http.ResponseWriter, r *http.Reques return } - // Simplified validation - in real implementation would use JSON schema validationType := r.URL.Query().Get("type") if validationType == "" { validationType = "input" } + if validationType != "input" && validationType != "output" { + http.Error(w, "type must be \"input\" or \"output\"", http.StatusBadRequest) + return + } + + // Validate against the schema this agent actually advertises. This + // previously answered "valid": true for every payload without inspecting + // it, so a client using the endpoint as a gate accepted anything. + schema := as.generateAgentSchema(agentID, as.agentMetaOrNil(agentID)) + errs := validateAgainstSchema(schemaSection(schema, validationType), data) + if errs == nil { + errs = []string{} + } response := map[string]interface{}{ - "valid": true, + "valid": len(errs) == 0, "agent_id": agentID, "type": validationType, - "errors": []string{}, + "errors": errs, "timestamp": time.Now().UTC().Format(time.RFC3339), } w.Header().Set("Content-Type", "application/json") - json.NewEncoder(w).Encode(response) + _ = json.NewEncoder(w).Encode(response) } // handleMetrics handles system metrics requests @@ -508,14 +541,14 @@ func (as *AutoServer) handleMetrics(w http.ResponseWriter, r *http.Request) { uptime := time.Since(as.startTime) metrics := map[string]interface{}{ - "total_agents": len(as.agentInstances), - "active_agents": len(as.agentInstances), - "requests": as.requestCount, + "total_agents": as.agentCount(), + "active_agents": as.agentCount(), + "requests": as.requestCount.Load(), "uptime": uptime.String(), "system": map[string]interface{}{ - "total_agents": len(as.agentInstances), - "active_agents": len(as.agentInstances), - "total_requests": as.requestCount, + "total_agents": as.agentCount(), + "active_agents": as.agentCount(), + "total_requests": as.requestCount.Load(), "uptime": uptime.String(), }, "agents": make(map[string]interface{}), @@ -523,7 +556,7 @@ func (as *AutoServer) handleMetrics(w http.ResponseWriter, r *http.Request) { } w.Header().Set("Content-Type", "application/json") - json.NewEncoder(w).Encode(metrics) + _ = json.NewEncoder(w).Encode(metrics) } // handleAgentMetrics handles agent-specific metrics requests @@ -531,7 +564,7 @@ func (as *AutoServer) handleAgentMetrics(w http.ResponseWriter, r *http.Request) vars := mux.Vars(r) agentID := vars["agentId"] - if _, exists := as.agentInstances[agentID]; !exists { + if _, exists := as.agentInstance(agentID); !exists { http.Error(w, "Agent not found", http.StatusNotFound) return } @@ -546,7 +579,7 @@ func (as *AutoServer) handleAgentMetrics(w http.ResponseWriter, r *http.Request) } w.Header().Set("Content-Type", "application/json") - json.NewEncoder(w).Encode(metrics) + _ = json.NewEncoder(w).Encode(metrics) } // Helper methods diff --git a/pkg/server/auto_security_test.go b/pkg/server/auto_security_test.go new file mode 100644 index 0000000..54333e6 --- /dev/null +++ b/pkg/server/auto_security_test.go @@ -0,0 +1,435 @@ +// Copyright (c) 2024 GoLangGraph Team +// +// Licensed under the MIT License. See LICENSE file in the project root for full license information. + +package server + +import ( + "bytes" + "context" + "encoding/json" + "net" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "strings" + "sync" + "testing" + "time" + + "github.com/UnicoLab/GoLangGraph/pkg/agent" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// The auto-generated server is a second, independent serving path: it is what +// the `auto-serve` CLI command runs. It previously had no authentication of any +// kind, a hardcoded wildcard CORS header, and a MaxRequestSize that was +// configured and never enforced β€” so hardening Server left all of this open. + +func newAutoServer(t *testing.T, mutate func(*AutoServerConfig)) *AutoServer { + t.Helper() + + cfg := DefaultAutoServerConfig() + cfg.Port = 0 + cfg.EnableWebUI = false + cfg.EnablePlayground = false + if mutate != nil { + mutate(cfg) + } + + // Each test gets its own registry: the default is process-wide, so tests + // sharing it would see each other's agents. + as := NewAutoServerWithRegistry(cfg, agent.NewAgentRegistry()) + + agentCfg := agent.DefaultAgentConfig() + agentCfg.ID = "auto-agent" + agentCfg.Name = "Auto Agent" + agentCfg.Type = agent.AgentTypeChat + agentCfg.Model = "fake-model" + agentCfg.Provider = "fake" + require.NoError(t, as.RegisterAgent("auto-agent", agent.NewBaseAgentDefinition(agentCfg))) + require.NoError(t, as.GenerateEndpoints()) + + return as +} + +func autoRequest(t *testing.T, as *AutoServer, method, path string, body string, headers map[string]string) *httptest.ResponseRecorder { + t.Helper() + + var reader *bytes.Reader + if body != "" { + reader = bytes.NewReader([]byte(body)) + } else { + reader = bytes.NewReader(nil) + } + + req := httptest.NewRequest(method, path, reader) + req.Header.Set("Content-Type", "application/json") + for k, v := range headers { + req.Header.Set(k, v) + } + rec := httptest.NewRecorder() + as.router.ServeHTTP(rec, req) + return rec +} + +// --------------------------------------------------------------------------- +// Authentication +// --------------------------------------------------------------------------- + +func TestAutoServer_AuthRequiredRejectsMissingKey(t *testing.T) { + as := newAutoServer(t, func(c *AutoServerConfig) { + c.Security.RequireAuth = true + c.Security.APIKeys = []string{"auto-key"} + }) + + rec := autoRequest(t, as, http.MethodGet, "/agents", "", nil) + assert.Equal(t, http.StatusUnauthorized, rec.Code, + "the auto server previously served every endpoint unauthenticated") + + rec = autoRequest(t, as, http.MethodGet, "/agents", "", map[string]string{"X-API-Key": "wrong"}) + assert.Equal(t, http.StatusUnauthorized, rec.Code) + + rec = autoRequest(t, as, http.MethodGet, "/agents", "", map[string]string{"X-API-Key": "auto-key"}) + assert.Equal(t, http.StatusOK, rec.Code) +} + +func TestAutoServer_HealthIsPublicUnderAuth(t *testing.T) { + as := newAutoServer(t, func(c *AutoServerConfig) { + c.Security.RequireAuth = true + c.Security.APIKeys = []string{"k"} + c.Security.PublicPaths = []string{"/health"} + }) + + rec := autoRequest(t, as, http.MethodGet, "/health", "", nil) + assert.Equal(t, http.StatusOK, rec.Code, "probes must not need credentials") +} + +func TestAutoServer_AuthDisabledByDefault(t *testing.T) { + as := newAutoServer(t, nil) + rec := autoRequest(t, as, http.MethodGet, "/agents", "", nil) + assert.Equal(t, http.StatusOK, rec.Code) +} + +func TestAutoServer_AuthWithNoKeysFailsClosed(t *testing.T) { + as := newAutoServer(t, func(c *AutoServerConfig) { + c.Security.RequireAuth = true + c.Security.APIKeys = nil + }) + rec := autoRequest(t, as, http.MethodGet, "/agents", "", map[string]string{"X-API-Key": "anything"}) + assert.Equal(t, http.StatusUnauthorized, rec.Code) +} + +// --------------------------------------------------------------------------- +// Cross-origin access +// --------------------------------------------------------------------------- + +func TestAutoServer_CORSRestrictsOrigins(t *testing.T) { + as := newAutoServer(t, func(c *AutoServerConfig) { + c.Security.AllowedOrigins = []string{"https://studio.example.com"} + }) + + rec := autoRequest(t, as, http.MethodGet, "/health", "", + map[string]string{"Origin": "https://studio.example.com"}) + assert.Equal(t, "https://studio.example.com", rec.Header().Get("Access-Control-Allow-Origin")) + assert.Contains(t, rec.Header().Get("Vary"), "Origin") + + rec = autoRequest(t, as, http.MethodGet, "/health", "", + map[string]string{"Origin": "https://evil.example.com"}) + assert.Empty(t, rec.Header().Get("Access-Control-Allow-Origin"), + "a hardcoded wildcard offered every origin access to every endpoint") +} + +// Most routes declare concrete methods, so preflight used to 404 and a browser +// blocked the request that followed. +func TestAutoServer_PreflightIsAnswered(t *testing.T) { + as := newAutoServer(t, nil) + + for _, path := range []string{"/agents", "/schemas", "/metrics", "/api/agents/auto-agent"} { + rec := autoRequest(t, as, http.MethodOptions, path, "", + map[string]string{"Origin": "http://localhost:3000"}) + assert.Less(t, rec.Code, 300, "preflight for %s must succeed, got %d", path, rec.Code) + assert.Contains(t, rec.Header().Get("Access-Control-Allow-Headers"), "X-API-Key", + "preflight for %s must permit the auth header", path) + } +} + +func TestAutoServer_PreflightRejectsDisallowedOrigin(t *testing.T) { + as := newAutoServer(t, func(c *AutoServerConfig) { + c.Security.AllowedOrigins = []string{"https://ok.example.com"} + }) + rec := autoRequest(t, as, http.MethodOptions, "/agents", "", + map[string]string{"Origin": "https://evil.example.com"}) + assert.Equal(t, http.StatusForbidden, rec.Code) +} + +// --------------------------------------------------------------------------- +// Request hygiene +// --------------------------------------------------------------------------- + +// MaxRequestSize was declared in the configuration and never enforced. +func TestAutoServer_EnforcesMaxRequestSize(t *testing.T) { + as := newAutoServer(t, func(c *AutoServerConfig) { + c.MaxRequestSize = 1024 + }) + + huge := `{"message":"` + strings.Repeat("a", 8192) + `"}` + rec := autoRequest(t, as, http.MethodPost, "/validate/auto-agent", huge, nil) + assert.NotEqual(t, http.StatusOK, rec.Code, "an oversized body must be rejected") +} + +func TestAutoServer_SecurityHeadersArePresent(t *testing.T) { + as := newAutoServer(t, nil) + rec := autoRequest(t, as, http.MethodGet, "/health", "", nil) + assert.Equal(t, "nosniff", rec.Header().Get("X-Content-Type-Options")) + assert.Equal(t, "DENY", rec.Header().Get("X-Frame-Options")) +} + +// Recovery was opt-in through a middleware name list, so a deployment that +// omitted it had a panicking handler tear down the connection. +func TestAutoServer_PanicIsRecoveredEvenWithoutOptIn(t *testing.T) { + as := newAutoServer(t, func(c *AutoServerConfig) { + c.Middleware = []string{"cors", "logging"} // deliberately no "recovery" + }) + as.router.HandleFunc("/boom", func(w http.ResponseWriter, r *http.Request) { + panic("handler exploded") + }) + + rec := autoRequest(t, as, http.MethodGet, "/boom", "", nil) + assert.Equal(t, http.StatusInternalServerError, rec.Code) + assert.NotContains(t, rec.Body.String(), "goroutine", "stack traces must not reach clients") +} + +func TestAutoServer_MalformedJSONIsRejected(t *testing.T) { + as := newAutoServer(t, nil) + + for _, body := range []string{"{", "", "null", "[]", `{"message":`} { + rec := autoRequest(t, as, http.MethodPost, "/validate/auto-agent", body, nil) + assert.NotEqual(t, http.StatusInternalServerError, rec.Code, + "malformed body %q must not produce a server error", body) + } +} + +// --------------------------------------------------------------------------- +// Schema validation on this serving path +// --------------------------------------------------------------------------- + +func TestAutoServer_ValidationActuallyValidates(t *testing.T) { + as := newAutoServer(t, nil) + + // A payload missing the required field must be rejected. This endpoint + // previously answered "valid": true for anything at all. + rec := autoRequest(t, as, http.MethodPost, "/validate/auto-agent", `{}`, nil) + require.Equal(t, http.StatusOK, rec.Code) + + var resp struct { + Valid bool `json:"valid"` + Errors []string `json:"errors"` + } + require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &resp)) + assert.False(t, resp.Valid, "a payload missing a required field must not validate") + assert.NotEmpty(t, resp.Errors) + + // A conforming payload validates. + rec = autoRequest(t, as, http.MethodPost, "/validate/auto-agent", `{"message":"hello"}`, nil) + require.Equal(t, http.StatusOK, rec.Code) + require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &resp)) + assert.True(t, resp.Valid, "a conforming payload must validate: %v", resp.Errors) +} + +func TestAutoServer_ValidationRejectsUnknownType(t *testing.T) { + as := newAutoServer(t, nil) + rec := autoRequest(t, as, http.MethodPost, "/validate/auto-agent?type=sideways", `{"message":"x"}`, nil) + assert.Equal(t, http.StatusBadRequest, rec.Code) +} + +func TestAutoServer_ValidationUnknownAgent(t *testing.T) { + as := newAutoServer(t, nil) + rec := autoRequest(t, as, http.MethodPost, "/validate/missing", `{"message":"x"}`, nil) + assert.Equal(t, http.StatusNotFound, rec.Code) +} + +// --------------------------------------------------------------------------- +// Concurrency +// --------------------------------------------------------------------------- + +// The metrics middleware increments a shared counter on every request. +func TestAutoServer_ConcurrentRequestsAreRaceFree(t *testing.T) { + as := newAutoServer(t, nil) + + var wg sync.WaitGroup + for i := 0; i < 16; i++ { + wg.Add(1) + go func() { + defer wg.Done() + for j := 0; j < 10; j++ { + autoRequest(t, as, http.MethodGet, "/health", "", nil) + autoRequest(t, as, http.MethodGet, "/metrics", "", nil) + } + }() + } + wg.Wait() +} + +// --------------------------------------------------------------------------- +// Lifecycle +// --------------------------------------------------------------------------- + +// Regenerating endpoints on a live server would rewrite the route table and the +// agent maps while request goroutines read them. +func TestAutoServer_RegenerateWhileRunningIsRefused(t *testing.T) { + as := newAutoServer(t, nil) + + as.started.Store(true) + err := as.GenerateEndpoints() + require.Error(t, err) + assert.Contains(t, err.Error(), "running") +} + +// Two servers must be able to hold different agents. The default registry is +// process-wide, which silently gave one server another's agents. +func TestAutoServer_IsolatedRegistriesDoNotShareAgents(t *testing.T) { + first := newAutoServer(t, nil) + second := newAutoServer(t, nil) + + extra := agent.DefaultAgentConfig() + extra.ID = "second-only" + extra.Name = "Second Only" + extra.Type = agent.AgentTypeChat + extra.Model = "fake-model" + extra.Provider = "fake" + require.NoError(t, second.RegisterAgent("second-only", agent.NewBaseAgentDefinition(extra))) + + _, existsInFirst := first.Registry().GetDefinition("second-only") + assert.False(t, existsInFirst, + "an agent registered on one server must not appear on another") + + _, existsInSecond := second.Registry().GetDefinition("second-only") + assert.True(t, existsInSecond) +} + +// Reading the agent maps concurrently with the metrics endpoint must be safe. +func TestAutoServer_AgentLookupsAreRaceFree(t *testing.T) { + as := newAutoServer(t, nil) + + var wg sync.WaitGroup + for i := 0; i < 16; i++ { + wg.Add(1) + go func() { + defer wg.Done() + for j := 0; j < 20; j++ { + _, _ = as.agentInstance("auto-agent") + _, _ = as.agentMeta("auto-agent") + _ = as.agentCount() + _ = as.agentIDs() + autoRequest(t, as, http.MethodGet, "/agents/auto-agent", "", nil) + } + }() + } + wg.Wait() +} + +// --------------------------------------------------------------------------- +// Start and directory loading +// --------------------------------------------------------------------------- + +// Start reported success when the port was already taken: ListenAndServe ran in +// a goroutine that only logged the failure, so Start blocked on ctx.Done() and +// the caller believed the server was up. +func TestAutoServer_StartReportsBindFailure(t *testing.T) { + occupied, err := net.Listen("tcp", "127.0.0.1:0") + require.NoError(t, err) + defer func() { _ = occupied.Close() }() + + port := occupied.Addr().(*net.TCPAddr).Port + + as := newAutoServer(t, func(c *AutoServerConfig) { + c.Host = "127.0.0.1" + c.Port = port + }) + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + startErr := make(chan error, 1) + go func() { startErr <- as.Start(ctx) }() + + select { + case err := <-startErr: + require.Error(t, err, "a taken port must be reported, not swallowed") + assert.Contains(t, err.Error(), "listen") + case <-time.After(5 * time.Second): + t.Fatal("Start blocked instead of reporting that the port was unavailable") + } +} + +// A server started on port 0 must report the port it actually got. +func TestAutoServer_AddressReportsBoundPort(t *testing.T) { + as := newAutoServer(t, func(c *AutoServerConfig) { + c.Host = "127.0.0.1" + c.Port = 0 + }) + + assert.Empty(t, as.Address(), "no address before Start") + + ctx, cancel := context.WithCancel(context.Background()) + done := make(chan error, 1) + go func() { done <- as.Start(ctx) }() + + require.Eventually(t, func() bool { return as.Address() != "" }, + 5*time.Second, 20*time.Millisecond, "Start must publish the bound address") + + assert.NotContains(t, as.Address(), ":0", "the reported port must be the real one") + + cancel() + select { + case <-done: + case <-time.After(10 * time.Second): + t.Fatal("Start did not return after its context was canceled") + } +} + +// LoadAgentsFromDirectory scanned nothing: it listed whatever was already +// registered and logged that count as though it had loaded them. +func TestAutoServer_LoadAgentsFromDirectory(t *testing.T) { + dir := t.TempDir() + + config := `name: from-directory +version: "1.0" +agents: + loaded-agent: + id: loaded-agent + name: Loaded Agent + type: chat + model: fake-model + provider: fake +` + require.NoError(t, os.WriteFile(filepath.Join(dir, "agents.yaml"), []byte(config), 0o600)) + + as := newAutoServer(t, nil) + require.NoError(t, as.LoadAgentsFromDirectory(dir)) + + _, exists := as.Registry().GetDefinition("loaded-agent") + assert.True(t, exists, "the agent defined in the directory must be registered") +} + +func TestAutoServer_LoadAgentsFromDirectoryReportsProblems(t *testing.T) { + as := newAutoServer(t, nil) + + // A directory that does not exist. + require.Error(t, as.LoadAgentsFromDirectory(filepath.Join(t.TempDir(), "missing"))) + + // A directory with nothing loadable must say so rather than report success. + empty := t.TempDir() + require.NoError(t, os.WriteFile(filepath.Join(empty, "notes.txt"), []byte("hi"), 0o600)) + err := as.LoadAgentsFromDirectory(empty) + require.Error(t, err) + assert.Contains(t, err.Error(), "no agent configuration files") + + // A malformed config must fail loudly. + bad := t.TempDir() + require.NoError(t, os.WriteFile(filepath.Join(bad, "broken.yaml"), []byte("{{{"), 0o600)) + assert.Error(t, as.LoadAgentsFromDirectory(bad)) +} diff --git a/pkg/server/auto_server.go b/pkg/server/auto_server.go index 26521c1..cfcba77 100644 --- a/pkg/server/auto_server.go +++ b/pkg/server/auto_server.go @@ -8,8 +8,16 @@ package server import ( "context" + "errors" "fmt" + "net" "net/http" + "os" + "path/filepath" + "sort" + "strings" + "sync" + "sync/atomic" "time" "github.com/gorilla/mux" @@ -34,8 +42,26 @@ type AutoServer struct { agentMetadata map[string]map[string]interface{} // Metrics tracking - startTime time.Time - requestCount int64 + startTime time.Time + // requestCount is incremented from every request goroutine, so it must be + // accessed atomically; a plain int64 here was a data race under any + // concurrent load, which is to say always. + requestCount atomic.Int64 + + // agentsMu guards agentInstances and agentMetadata. GenerateEndpoints + // writes them and every handler reads them, and the public API permits + // registering an agent and regenerating after the server is serving. + agentsMu sync.RWMutex + + // started records whether Start has been called, so endpoints cannot be + // regenerated underneath live traffic. + started atomic.Bool + + // mu guards boundAddr. + mu sync.Mutex + // boundAddr is the address actually listened on, which differs from the + // configured one when port 0 is used. + boundAddr string } // AutoServerConfig configures the auto-generated server @@ -54,6 +80,11 @@ type AutoServerConfig struct { ServerTimeout time.Duration `yaml:"server_timeout" json:"server_timeout"` MaxRequestSize int64 `yaml:"max_request_size" json:"max_request_size"` Middleware []string `yaml:"middleware" json:"middleware"` + + // Security controls authentication, allowed origins and request limits, + // using the same configuration type as Server. Nil falls back to + // DefaultSecurityConfig. + Security *SecurityConfig `yaml:"security" json:"security"` } // DefaultAutoServerConfig returns default configuration @@ -72,10 +103,17 @@ func DefaultAutoServerConfig() *AutoServerConfig { ServerTimeout: 30 * time.Second, MaxRequestSize: 10 * 1024 * 1024, // 10MB Middleware: []string{"cors", "logging", "recovery"}, + Security: DefaultSecurityConfig(), } } -// NewAutoServer creates a new auto-server instance +// NewAutoServer creates a new auto-server instance backed by the process-wide +// agent registry. +// +// Note that the registry is shared: two AutoServer instances in one process see +// each other's agents, so an agent registered for one is served by the other. +// That is rarely what you want when the two servers have different exposure or +// credentials. Use NewAutoServerWithRegistry to give a server its own registry. func NewAutoServer(config *AutoServerConfig) *AutoServer { if config == nil { config = DefaultAutoServerConfig() @@ -101,19 +139,52 @@ func NewAutoServer(config *AutoServerConfig) *AutoServer { agentInstances: make(map[string]agent.Agent), agentMetadata: make(map[string]map[string]interface{}), startTime: time.Now(), - requestCount: 0, } } // LoadAgentsFromDirectory loads agent definitions from a directory +// LoadAgentsFromDirectory loads every agent configuration file in a directory. +// +// This previously scanned nothing: it listed whatever was already registered +// and logged that count as "Loaded agent definitions", so a caller pointing at +// a directory of configs got silence and no agents. func (as *AutoServer) LoadAgentsFromDirectory(directory string) error { as.logger.WithField("directory", directory).Info("Loading agents from directory") - // This would scan for Go files and load agent definitions - // For now, we'll use the existing registry system - definitions := as.registry.ListDefinitions() + entries, err := os.ReadDir(directory) + if err != nil { + return fmt.Errorf("failed to read agent directory %s: %w", directory, err) + } + + names := make([]string, 0, len(entries)) + for _, entry := range entries { + if entry.IsDir() { + continue + } + switch strings.ToLower(filepath.Ext(entry.Name())) { + case ".yaml", ".yml", ".json": + names = append(names, entry.Name()) + } + } + sort.Strings(names) + + if len(names) == 0 { + return fmt.Errorf("no agent configuration files (.yaml, .yml, .json) found in %s", directory) + } + + loaded := 0 + for _, name := range names { + path := filepath.Join(directory, name) + if err := as.LoadAgentsFromConfig(path); err != nil { + return fmt.Errorf("failed to load %s: %w", path, err) + } + loaded++ + } - as.logger.WithField("count", len(definitions)).Info("Loaded agent definitions") + as.logger.WithFields(logrus.Fields{ + "directory": directory, + "files": loaded, + }).Info("Loaded agent definitions from directory") return nil } @@ -136,6 +207,70 @@ func (as *AutoServer) LoadAgentsFromConfig(configPath string) error { return nil } +// NewAutoServerWithRegistry creates an auto-server with its own agent registry, +// isolated from the process-wide one and from any other server. +func NewAutoServerWithRegistry(config *AutoServerConfig, registry *agent.AgentRegistry) *AutoServer { + as := NewAutoServer(config) + if registry != nil { + as.registry = registry + } + return as +} + +// agentInstance returns a registered agent instance. +func (as *AutoServer) agentInstance(id string) (agent.Agent, bool) { + as.agentsMu.RLock() + defer as.agentsMu.RUnlock() + instance, ok := as.agentInstances[id] + return instance, ok +} + +// agentMeta returns a registered agent's metadata. +func (as *AutoServer) agentMeta(id string) (map[string]interface{}, bool) { + as.agentsMu.RLock() + defer as.agentsMu.RUnlock() + metadata, ok := as.agentMetadata[id] + return metadata, ok +} + +// agentMetaOrNil returns an agent's metadata, or nil when it is not registered. +func (as *AutoServer) agentMetaOrNil(id string) map[string]interface{} { + metadata, _ := as.agentMeta(id) + return metadata +} + +// agentCount returns how many agents are being served. +func (as *AutoServer) agentCount() int { + as.agentsMu.RLock() + defer as.agentsMu.RUnlock() + return len(as.agentInstances) +} + +// agentIDs returns the served agent IDs. +func (as *AutoServer) agentIDs() []string { + as.agentsMu.RLock() + defer as.agentsMu.RUnlock() + ids := make([]string, 0, len(as.agentInstances)) + for id := range as.agentInstances { + ids = append(ids, id) + } + sort.Strings(ids) + return ids +} + +// Address returns the address the server is listening on, or an empty string +// before Start. With port 0 configured this reports the port actually chosen. +func (as *AutoServer) Address() string { + as.mu.Lock() + defer as.mu.Unlock() + return as.boundAddr +} + +// Registry returns the agent registry this server serves from. +func (as *AutoServer) Registry() *agent.AgentRegistry { + return as.registry +} + // RegisterAgent registers a single agent programmatically func (as *AutoServer) RegisterAgent(id string, definition agent.AgentDefinition) error { return as.registry.RegisterDefinition(id, definition) @@ -143,6 +278,12 @@ func (as *AutoServer) RegisterAgent(id string, definition agent.AgentDefinition) // GenerateEndpoints automatically generates REST endpoints for all registered agents func (as *AutoServer) GenerateEndpoints() error { + // Regenerating on a live server would mutate the route table and the agent + // maps while handlers are reading them. + if as.started.Load() { + return fmt.Errorf("cannot generate endpoints while the server is running") + } + as.logger.Info("Generating dynamic endpoints for agents") // Apply middleware @@ -171,6 +312,14 @@ func (as *AutoServer) GenerateEndpoints() error { as.generateMetricsEndpoints() } + // Cross-origin preflight. Most routes declare concrete methods, so an + // OPTIONS request fell through to 404 and a browser blocked every + // cross-origin POST, PUT and DELETE against them. Registered last so it + // only catches what no other route matched. + as.router.PathPrefix("/").Methods(http.MethodOptions).HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusNoContent) + }) + as.logger.Info("Successfully generated all endpoints") return nil } @@ -209,8 +358,10 @@ func (as *AutoServer) generateAgentEndpoints() error { continue } + as.agentsMu.Lock() as.agentInstances[agentID] = agentInstance as.agentMetadata[agentID] = definition.GetMetadata() + as.agentsMu.Unlock() // Generate endpoints for this agent basePath := fmt.Sprintf("%s/%s", as.config.BasePath, agentID) @@ -277,23 +428,124 @@ func (as *AutoServer) generateMetricsEndpoints() { as.logger.Info("Generated metrics endpoints") } -// applyMiddleware applies configured middleware +// applyMiddleware applies the middleware chain. +// +// Recovery, the request size limit and the security headers are unconditional: +// they were previously opt-in through config.Middleware, so a deployment that +// omitted "recovery" from that list had a panicking handler tear down the +// connection, and MaxRequestSize was configured but never enforced at all. +// Recovery is registered first so it wraps everything that follows. func (as *AutoServer) applyMiddleware() { - // Always apply metrics middleware + if as.config.Security == nil { + as.config.Security = DefaultSecurityConfig() + } + // A request limit set on the server config wins over the security default. + if as.config.MaxRequestSize > 0 { + as.config.Security.MaxRequestBytes = as.config.MaxRequestSize + } + + as.router.Use(recoveryMiddleware(as.logger)) + as.router.Use(bodyLimitMiddleware(as.config.Security.maxBytes())) + as.router.Use(securityHeadersMiddleware) as.router.Use(as.metricsMiddleware()) for _, middleware := range as.config.Middleware { switch middleware { case "cors": if as.config.EnableCORS { - as.router.Use(corsMiddleware) + as.router.Use(as.corsMiddleware()) } case "logging": as.router.Use(loggingMiddleware(as.logger)) case "recovery": - as.router.Use(recoveryMiddleware(as.logger)) + // Applied unconditionally above. } } + + // Authentication runs innermost so rejected requests still carry the CORS + // and security headers a browser needs to read the response. + as.router.Use(as.authMiddleware()) +} + +// corsMiddleware answers cross-origin requests against the configured origin +// allowlist. It previously emitted a hardcoded "*" with no way to restrict it. +func (as *AutoServer) corsMiddleware() func(http.Handler) http.Handler { + return func(next http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + origin := r.Header.Get("Origin") + w.Header().Add("Vary", "Origin") + + if !as.config.Security.originAllowed(origin) { + if r.Method == http.MethodOptions { + w.WriteHeader(http.StatusForbidden) + return + } + next.ServeHTTP(w, r) + return + } + + if allow := as.config.Security.corsOrigin(origin); allow != "" { + w.Header().Set("Access-Control-Allow-Origin", allow) + if allow != "*" { + w.Header().Set("Access-Control-Allow-Credentials", "true") + } + } + w.Header().Set("Access-Control-Allow-Methods", "GET, POST, PUT, DELETE, OPTIONS") + w.Header().Set("Access-Control-Allow-Headers", "Content-Type, Authorization, X-API-Key") + w.Header().Set("Access-Control-Max-Age", "600") + + if r.Method == http.MethodOptions { + w.WriteHeader(http.StatusOK) + return + } + + next.ServeHTTP(w, r) + }) + } +} + +// authMiddleware enforces API key authentication when configured. The +// auto-generated server previously had no authentication of any kind. +func (as *AutoServer) authMiddleware() func(http.Handler) http.Handler { + return func(next http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + sec := as.config.Security + if sec == nil || !sec.RequireAuth { + next.ServeHTTP(w, r) + return + } + if r.Method == http.MethodOptions || sec.isPublic(r.URL.Path) { + next.ServeHTTP(w, r) + return + } + + if !sec.authorized(r.Header.Get("X-API-Key")) { + as.logger.WithFields(logrus.Fields{ + "path": r.URL.Path, + "remote": r.RemoteAddr, + }).Warn("Rejected unauthenticated request") + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusUnauthorized) + _, _ = w.Write([]byte(`{"error":"missing or invalid API key"}`)) + return + } + + next.ServeHTTP(w, r) + }) + } +} + +// allowedOrigin returns the value a streaming handler should echo, so +// server-sent event responses honor the same allowlist as everything else. +func (as *AutoServer) allowedOrigin(r *http.Request) string { + if as.config == nil || as.config.Security == nil { + return "" + } + origin := r.Header.Get("Origin") + if !as.config.Security.originAllowed(origin) { + return "" + } + return as.config.Security.corsOrigin(origin) } // Start starts the auto-server @@ -301,6 +553,9 @@ func (as *AutoServer) Start(ctx context.Context) error { if err := as.GenerateEndpoints(); err != nil { return fmt.Errorf("failed to generate endpoints: %w", err) } + // From here on the route table and agent maps are read by request + // goroutines and must not be regenerated. + as.started.Store(true) address := fmt.Sprintf("%s:%d", as.config.Host, as.config.Port) @@ -311,19 +566,41 @@ func (as *AutoServer) Start(ctx context.Context) error { WriteTimeout: as.config.ServerTimeout, } - as.logger.WithField("address", address).Info("Starting auto-generated multi-agent server") + // Bind before reporting success. ListenAndServe used to run in a goroutine + // that merely logged a failure, so Start blocked on ctx.Done() and the + // caller believed the server was up when the port was already taken. + listener, err := net.Listen("tcp", address) + if err != nil { + return fmt.Errorf("failed to listen on %s: %w", address, err) + } + + as.mu.Lock() + as.boundAddr = listener.Addr().String() + as.mu.Unlock() + + as.logger.WithField("address", listener.Addr().String()).Info("Starting auto-generated multi-agent server") // Print available endpoints as.printEndpoints() + serveErr := make(chan error, 1) go func() { - if err := server.ListenAndServe(); err != nil && err != http.ErrServerClosed { - as.logger.WithError(err).Error("Server failed to start") + if err := server.Serve(listener); err != nil && !errors.Is(err, http.ErrServerClosed) { + as.logger.WithError(err).Error("Server stopped unexpectedly") + serveErr <- err + return } + serveErr <- nil }() - // Wait for shutdown signal - <-ctx.Done() + // Wait for shutdown, or for the server to stop on its own. + select { + case <-ctx.Done(): + case err := <-serveErr: + if err != nil { + return fmt.Errorf("server stopped: %w", err) + } + } as.logger.Info("Shutting down server...") shutdownCtx, cancel := context.WithTimeout(context.Background(), 30*time.Second) @@ -386,7 +663,7 @@ func setupLLMProviders(manager *llm.ProviderManager, config *AutoServerConfig) { // Just skip this provider if it fails return } - manager.RegisterProvider("ollama", ollamaProvider) + _ = manager.RegisterProvider("ollama", ollamaProvider) } // Setup other providers from config @@ -402,27 +679,13 @@ func setupLLMProviders(manager *llm.ProviderManager, config *AutoServerConfig) { func (as *AutoServer) metricsMiddleware() func(http.Handler) http.Handler { return func(next http.Handler) http.Handler { return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - as.requestCount++ + as.requestCount.Add(1) next.ServeHTTP(w, r) }) } } // Middleware functions -func corsMiddleware(next http.Handler) http.Handler { - return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - w.Header().Set("Access-Control-Allow-Origin", "*") - w.Header().Set("Access-Control-Allow-Methods", "GET, POST, PUT, DELETE, OPTIONS") - w.Header().Set("Access-Control-Allow-Headers", "Content-Type, Authorization") - - if r.Method == "OPTIONS" { - w.WriteHeader(http.StatusOK) - return - } - - next.ServeHTTP(w, r) - }) -} func loggingMiddleware(logger *logrus.Logger) func(http.Handler) http.Handler { return func(next http.Handler) http.Handler { @@ -438,17 +701,3 @@ func loggingMiddleware(logger *logrus.Logger) func(http.Handler) http.Handler { }) } } - -func recoveryMiddleware(logger *logrus.Logger) func(http.Handler) http.Handler { - return func(next http.Handler) http.Handler { - return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - defer func() { - if err := recover(); err != nil { - logger.WithField("error", err).Error("Panic recovered") - http.Error(w, "Internal Server Error", http.StatusInternalServerError) - } - }() - next.ServeHTTP(w, r) - }) - } -} diff --git a/pkg/server/frontend_contract_test.go b/pkg/server/frontend_contract_test.go index c535268..4127f08 100644 --- a/pkg/server/frontend_contract_test.go +++ b/pkg/server/frontend_contract_test.go @@ -89,7 +89,7 @@ func TestFrontendAPIContract(t *testing.T) { t.Errorf("expected agent name, got %v", get.Agent["name"]) } - // 3. Execute agent -> {"execution": {PascalCase...}} + // 3. Execute agent -> {"execution": {...snake_case...}} execRaw := doRequest(t, server, "POST", "/api/v1/agents/"+agentID+"/execute", `{"input":"hello"}`) var exec struct { Execution map[string]interface{} `json:"execution"` @@ -100,13 +100,29 @@ func TestFrontendAPIContract(t *testing.T) { if exec.Execution == nil { t.Fatal("expected an execution object in execute response") } - for _, key := range []string{"ID", "Input", "Output", "Success", "Status", "Duration", "Steps", "ToolCalls"} { + // AgentExecution used to ship untagged, so it alone on this API serialized + // as Go PascalCase while every neighbouring payload was snake_case β€” and + // its Error field, being a Go error, marshalled to {} so a failed run + // reached the client with no reason in it. Both are fixed by tagging the + // struct; this pins the tagged names so they cannot silently regress. + for _, key := range []string{"id", "input", "output", "success", "status", "duration", "tool_calls", "execution_path", "timestamp"} { if _, ok := exec.Execution[key]; !ok { - t.Errorf("execution missing PascalCase field %q", key) + t.Errorf("execution missing snake_case field %q", key) + } + } + for _, key := range []string{"ID", "Input", "Output", "Success", "Status", "Duration", "Steps", "ToolCalls"} { + if _, ok := exec.Execution[key]; ok { + t.Errorf("execution still exposes untagged Go field %q", key) } } - if success, _ := exec.Execution["Success"].(bool); !success { - t.Errorf("expected successful execution, got %v", exec.Execution["Success"]) + if success, _ := exec.Execution["success"].(bool); !success { + t.Errorf("expected successful execution, got %v", exec.Execution["success"]) + } + // A Go error is not serialisable; the reason must travel as a string. + if raw, ok := exec.Execution["error"]; ok { + if _, isString := raw.(string); !isString { + t.Errorf("execution error must serialize as a string, got %T", raw) + } } } diff --git a/pkg/server/fuzz_test.go b/pkg/server/fuzz_test.go new file mode 100644 index 0000000..1703144 --- /dev/null +++ b/pkg/server/fuzz_test.go @@ -0,0 +1,122 @@ +// Copyright (c) 2024 GoLangGraph Team +// +// Licensed under the MIT License. See LICENSE file in the project root for full license information. + +package server + +import ( + "net/http" + "net/http/httptest" + "net/url" + "strings" + "testing" +) + +// FuzzAPIRequestBodies feeds arbitrary payloads to the write endpoints. Request +// bodies come from the network, so a malformed one must produce a 4xx, never a +// panic and never a 500 from an unhandled decode. +func FuzzAPIRequestBodies(f *testing.F) { + seeds := []string{ + `{"input":"hello"}`, + `{"state":{"a":1}}`, + `{"input":"x","thread_id":"t"}`, + `{}`, `[]`, `null`, `true`, ``, `{`, + `{"input":12345}`, + `{"state":"not-an-object"}`, + `{"state":{"deep":{"deep":{"deep":[1,2,3]}}}}`, + strings.Repeat(`{"a":`, 500), + "\x00\x01\x02", + } + for _, s := range seeds { + f.Add(s) + } + + paths := []string{ + "/api/v1/graphs/fuzz/execute", + "/api/v1/graphs/missing/execute", + "/api/v1/agents", + "/api/v1/sessions", + "/api/v1/threads", + } + + f.Fuzz(func(t *testing.T, body string) { + s := fuzzServer(t) + + for _, path := range paths { + func() { + defer func() { + if r := recover(); r != nil { + t.Fatalf("handler for %s panicked on %q: %v", path, body, r) + } + }() + + req := httptest.NewRequest(http.MethodPost, path, strings.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + rec := httptest.NewRecorder() + s.router.ServeHTTP(rec, req) + + if rec.Code == http.StatusInternalServerError { + t.Fatalf("%s returned 500 for body %q: %s", path, body, rec.Body.String()) + } + }() + } + }) +} + +// FuzzAPIPaths checks routing against arbitrary path segments: an identifier +// from a URL must never panic a handler or escape into a server error. +func FuzzAPIPaths(f *testing.F) { + seeds := []string{ + "fuzz", "missing", "", "..", "../../etc/passwd", "%2e%2e", + strings.Repeat("a", 4096), "with space", "unicode-Γ©", + } + for _, s := range seeds { + f.Add(s) + } + + f.Fuzz(func(t *testing.T, id string) { + s := fuzzServer(t) + + for _, tmpl := range []string{ + "/api/v1/graphs/%s", + "/api/v1/graphs/%s/topology", + "/api/v1/agents/%s", + "/api/v1/tools/%s", + } { + func() { + defer func() { + if r := recover(); r != nil { + t.Fatalf("path %q panicked with id %q: %v", tmpl, id, r) + } + }() + + req := httptest.NewRequest(http.MethodGet, safePath(tmpl, id), nil) + rec := httptest.NewRecorder() + s.router.ServeHTTP(rec, req) + + if rec.Code == http.StatusInternalServerError { + t.Fatalf("%q with id %q returned 500: %s", tmpl, id, rec.Body.String()) + } + }() + } + }) +} + +// safePath substitutes an identifier into a path template, percent-encoding it +// the way a real HTTP client would. Without this, httptest.NewRequest panics +// while *building* the request, which would report a harness failure rather +// than a server one. +func safePath(tmpl, id string) string { + return strings.Replace(tmpl, "%s", url.PathEscape(id), 1) +} + +// fuzzServer builds a server with one registered graph, without the static +// catch-all so route resolution is exercised directly. +func fuzzServer(t *testing.T) *Server { + t.Helper() + s := newTestServer(t, nil) + if g, ok := s.GraphManager().Get("demo"); ok { + s.GraphManager().Register("fuzz", g) + } + return s +} diff --git a/pkg/server/graph_api.go b/pkg/server/graph_api.go new file mode 100644 index 0000000..2c0ecc4 --- /dev/null +++ b/pkg/server/graph_api.go @@ -0,0 +1,238 @@ +// Copyright (c) 2024 GoLangGraph Team +// +// Licensed under the MIT License. See LICENSE file in the project root for full license information. +// +// Package: GoLangGraph - A powerful Go framework for building AI agent workflows + +package server + +import ( + "sort" + "sync" + + "github.com/UnicoLab/GoLangGraph/pkg/core" +) + +// GraphManager holds the graphs a server exposes over the API. Registering a +// graph makes it listable, inspectable, executable and streamable, which is +// what GoLangGraph Studio needs to debug a workflow. +type GraphManager struct { + mu sync.RWMutex + graphs map[string]*core.Graph + order []string +} + +// NewGraphManager creates an empty graph manager. +func NewGraphManager() *GraphManager { + return &GraphManager{graphs: make(map[string]*core.Graph)} +} + +// Register adds a graph under an ID. Re-registering an ID replaces the graph. +func (gm *GraphManager) Register(id string, g *core.Graph) { + if gm == nil || g == nil || id == "" { + return + } + gm.mu.Lock() + defer gm.mu.Unlock() + if _, exists := gm.graphs[id]; !exists { + gm.order = append(gm.order, id) + } + gm.graphs[id] = g +} + +// Unregister removes a graph. +func (gm *GraphManager) Unregister(id string) { + if gm == nil { + return + } + gm.mu.Lock() + defer gm.mu.Unlock() + delete(gm.graphs, id) + for i, existing := range gm.order { + if existing == id { + gm.order = append(gm.order[:i], gm.order[i+1:]...) + break + } + } +} + +// Get returns a graph by ID. +func (gm *GraphManager) Get(id string) (*core.Graph, bool) { + if gm == nil { + return nil, false + } + gm.mu.RLock() + defer gm.mu.RUnlock() + g, ok := gm.graphs[id] + return g, ok +} + +// List returns registered graph IDs in registration order. +func (gm *GraphManager) List() []string { + if gm == nil { + return nil + } + gm.mu.RLock() + defer gm.mu.RUnlock() + return append([]string(nil), gm.order...) +} + +// GraphNodeView describes a node for API clients. +type GraphNodeView struct { + ID string `json:"id"` + Name string `json:"name"` + Type string `json:"type"` + IsStart bool `json:"is_start"` + IsEnd bool `json:"is_end"` + Metadata map[string]interface{} `json:"metadata,omitempty"` +} + +// GraphEdgeView describes an edge for API clients. +type GraphEdgeView struct { + From string `json:"from"` + To string `json:"to"` + Conditional bool `json:"conditional"` + // RouteKey is set for conditional edges and names the routing key that + // selects this destination. + RouteKey string `json:"route_key,omitempty"` +} + +// GraphTopologyView is the serialisable topology of a graph. Studio renders +// this directly, so both nodes and edges are always present (never null). +type GraphTopologyView struct { + Nodes []GraphNodeView `json:"nodes"` + Edges []GraphEdgeView `json:"edges"` +} + +// GraphSummaryView describes a graph without its topology. +type GraphSummaryView struct { + ID string `json:"id"` + Name string `json:"name"` + StartNode string `json:"start_node"` + EndNodes []string `json:"end_nodes"` + NodeCount int `json:"node_count"` + EdgeCount int `json:"edge_count"` + Running bool `json:"running"` +} + +// summariseGraph builds a summary view of a graph. +func summariseGraph(id string, g *core.Graph) GraphSummaryView { + topo := describeTopology(g) + endNodes := append([]string(nil), g.EndNodes...) + if endNodes == nil { + endNodes = []string{} + } + return GraphSummaryView{ + ID: id, + Name: g.Name, + StartNode: g.StartNode, + EndNodes: endNodes, + NodeCount: len(topo.Nodes), + EdgeCount: len(topo.Edges), + Running: g.IsRunning(), + } +} + +// describeTopology converts a graph into its serialisable topology, including +// conditional routes so a client sees every reachable path. +func describeTopology(g *core.Graph) GraphTopologyView { + view := GraphTopologyView{Nodes: []GraphNodeView{}, Edges: []GraphEdgeView{}} + if g == nil { + return view + } + + ids := make([]string, 0, len(g.Nodes)) + for id := range g.Nodes { + ids = append(ids, id) + } + sort.Strings(ids) + + for _, id := range ids { + node := g.Nodes[id] + if node == nil { + continue + } + nodeType := "node" + if t, ok := node.Metadata["type"].(string); ok && t != "" { + nodeType = t + } + view.Nodes = append(view.Nodes, GraphNodeView{ + ID: node.ID, + Name: node.Name, + Type: nodeType, + IsStart: node.ID == g.StartNode, + IsEnd: g.IsEndNode(node.ID), + Metadata: node.Metadata, + }) + } + + // Static edges, ordered for a stable rendering. + static := make([]GraphEdgeView, 0, len(g.Edges)) + for _, edge := range g.Edges { + if edge == nil { + continue + } + static = append(static, GraphEdgeView{ + From: edge.From, + To: edge.To, + Conditional: edge.Condition != nil, + }) + } + sort.Slice(static, func(i, j int) bool { + if static[i].From != static[j].From { + return static[i].From < static[j].From + } + return static[i].To < static[j].To + }) + view.Edges = append(view.Edges, static...) + + // Conditional routes registered via AddConditionalEdges. + for _, id := range ids { + ce, ok := g.GetConditionalEdge(id) + if !ok || ce == nil { + continue + } + keys := make([]string, 0, len(ce.Routes)) + for k := range ce.Routes { + keys = append(keys, k) + } + sort.Strings(keys) + for _, key := range keys { + view.Edges = append(view.Edges, GraphEdgeView{ + From: id, + To: ce.Routes[key], + Conditional: true, + RouteKey: key, + }) + } + } + + return view +} + +// ExecutionStepView is a single node execution, as sent to clients. +type ExecutionStepView struct { + NodeID string `json:"node_id"` + Step int `json:"step"` + Success bool `json:"success"` + Error string `json:"error,omitempty"` + DurationM float64 `json:"duration_ms"` + Attempts int `json:"attempts"` + State map[string]core.StateValue `json:"state,omitempty"` +} + +// describeStep converts an engine result into its wire representation. +func describeStep(r *core.ExecutionResult) ExecutionStepView { + view := ExecutionStepView{ + NodeID: r.NodeID, + Step: r.Step, + Success: r.Success, + Error: r.ErrorMessage, + DurationM: float64(r.Duration.Microseconds()) / 1000.0, + Attempts: r.Attempts, + } + if r.State != nil { + view.State = r.State.GetAll() + } + return view +} diff --git a/pkg/server/schema_validation.go b/pkg/server/schema_validation.go new file mode 100644 index 0000000..39b8a64 --- /dev/null +++ b/pkg/server/schema_validation.go @@ -0,0 +1,214 @@ +// Copyright (c) 2024 GoLangGraph Team +// +// Licensed under the MIT License. See LICENSE file in the project root for full license information. +// +// Package: GoLangGraph - A powerful Go framework for building AI agent workflows + +package server + +import ( + "fmt" + "sort" + "strings" +) + +// validateAgainstSchema checks a decoded JSON value against the JSON Schema +// subset that generateAgentSchema produces: type, properties, required, +// additionalProperties, minLength, maxLength, minimum, maximum, enum and items. +// +// The validation endpoint previously answered "valid": true for every input +// without inspecting it, so a client using it as a gate accepted anything. +// Unsupported keywords are ignored rather than treated as failures, so a richer +// schema is never reported invalid for a rule this validator cannot check. +func validateAgainstSchema(schema map[string]interface{}, value interface{}) []string { + if len(schema) == 0 { + return nil + } + var errs []string + validateValue(schema, value, "", &errs) + sort.Strings(errs) + return errs +} + +func fieldLabel(path string) string { + if path == "" { + return "value" + } + return path +} + +func join(path, key string) string { + if path == "" { + return key + } + return path + "." + key +} + +func validateValue(schema map[string]interface{}, value interface{}, path string, errs *[]string) { + expected, _ := schema["type"].(string) + + if enum, ok := schema["enum"].([]interface{}); ok && len(enum) > 0 { + matched := false + for _, candidate := range enum { + if fmt.Sprint(candidate) == fmt.Sprint(value) { + matched = true + break + } + } + if !matched { + *errs = append(*errs, fmt.Sprintf("%s must be one of %v", fieldLabel(path), enum)) + } + } + + switch expected { + case "object": + obj, ok := value.(map[string]interface{}) + if !ok { + *errs = append(*errs, fmt.Sprintf("%s must be an object", fieldLabel(path))) + return + } + validateObject(schema, obj, path, errs) + + case "array": + items, ok := toArray(value) + if !ok { + *errs = append(*errs, fmt.Sprintf("%s must be an array", fieldLabel(path))) + return + } + if itemSchema, ok := schema["items"].(map[string]interface{}); ok { + for i, item := range items { + validateValue(itemSchema, item, fmt.Sprintf("%s[%d]", fieldLabel(path), i), errs) + } + } + + case "string": + str, ok := value.(string) + if !ok { + *errs = append(*errs, fmt.Sprintf("%s must be a string", fieldLabel(path))) + return + } + if min, ok := toNumber(schema["minLength"]); ok && float64(len(str)) < min { + *errs = append(*errs, fmt.Sprintf("%s must be at least %d characters", fieldLabel(path), int(min))) + } + if max, ok := toNumber(schema["maxLength"]); ok && float64(len(str)) > max { + *errs = append(*errs, fmt.Sprintf("%s must be at most %d characters", fieldLabel(path), int(max))) + } + + case "number", "integer": + num, ok := toNumber(value) + if !ok { + *errs = append(*errs, fmt.Sprintf("%s must be a number", fieldLabel(path))) + return + } + if expected == "integer" && num != float64(int64(num)) { + *errs = append(*errs, fmt.Sprintf("%s must be an integer", fieldLabel(path))) + } + if min, ok := toNumber(schema["minimum"]); ok && num < min { + *errs = append(*errs, fmt.Sprintf("%s must be at least %v", fieldLabel(path), min)) + } + if max, ok := toNumber(schema["maximum"]); ok && num > max { + *errs = append(*errs, fmt.Sprintf("%s must be at most %v", fieldLabel(path), max)) + } + + case "boolean": + if _, ok := value.(bool); !ok { + *errs = append(*errs, fmt.Sprintf("%s must be a boolean", fieldLabel(path))) + } + + case "": + // No declared type: only structural keywords apply. + if obj, ok := value.(map[string]interface{}); ok { + validateObject(schema, obj, path, errs) + } + } +} + +func validateObject(schema map[string]interface{}, obj map[string]interface{}, path string, errs *[]string) { + properties, _ := schema["properties"].(map[string]interface{}) + + if required, ok := toStringSlice(schema["required"]); ok { + for _, key := range required { + if _, present := obj[key]; !present { + *errs = append(*errs, fmt.Sprintf("%s is required", join(path, key))) + } + } + } + + if len(properties) > 0 { + keys := make([]string, 0, len(obj)) + for k := range obj { + keys = append(keys, k) + } + sort.Strings(keys) + + for _, key := range keys { + propSchema, declared := properties[key].(map[string]interface{}) + if !declared { + if allow, ok := schema["additionalProperties"].(bool); ok && !allow { + *errs = append(*errs, fmt.Sprintf("%s is not an allowed property", join(path, key))) + } + continue + } + validateValue(propSchema, obj[key], join(path, key), errs) + } + } +} + +func toArray(value interface{}) ([]interface{}, bool) { + switch v := value.(type) { + case []interface{}: + return v, true + case nil: + return nil, false + } + return nil, false +} + +func toNumber(value interface{}) (float64, bool) { + switch v := value.(type) { + case float64: + return v, true + case float32: + return float64(v), true + case int: + return float64(v), true + case int32: + return float64(v), true + case int64: + return float64(v), true + } + return 0, false +} + +func toStringSlice(value interface{}) ([]string, bool) { + switch v := value.(type) { + case []string: + return v, true + case []interface{}: + out := make([]string, 0, len(v)) + for _, item := range v { + out = append(out, fmt.Sprint(item)) + } + return out, true + } + return nil, false +} + +// schemaSection extracts the "input" or "output" half of a generated schema. +func schemaSection(schema map[string]interface{}, section string) map[string]interface{} { + if schema == nil { + return nil + } + if sub, ok := schema[section].(map[string]interface{}); ok { + return sub + } + return nil +} + +// summariseErrors renders validation errors for a log line. +func summariseErrors(errs []string) string { + if len(errs) == 0 { + return "" + } + return strings.Join(errs, "; ") +} diff --git a/pkg/server/schema_validation_test.go b/pkg/server/schema_validation_test.go new file mode 100644 index 0000000..40334e2 --- /dev/null +++ b/pkg/server/schema_validation_test.go @@ -0,0 +1,183 @@ +// Copyright (c) 2024 GoLangGraph Team +// +// Licensed under the MIT License. See LICENSE file in the project root for full license information. + +package server + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// The schema the auto server advertises for an agent's input. +func inputSchema() map[string]interface{} { + return map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "message": map[string]interface{}{ + "type": "string", + "minLength": 1, + "maxLength": 10, + }, + "count": map[string]interface{}{ + "type": "integer", + "minimum": 0, + "maximum": 5, + }, + "mode": map[string]interface{}{ + "type": "string", + "enum": []interface{}{"fast", "slow"}, + }, + "history": map[string]interface{}{ + "type": "array", + "items": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "role": map[string]interface{}{"type": "string"}, + }, + "required": []string{"role"}, + }, + }, + }, + "required": []string{"message"}, + } +} + +// A conforming payload validates. The endpoint previously said this for every +// payload, so the interesting cases are the ones below it. +func TestSchemaValidation_AcceptsValidInput(t *testing.T) { + errs := validateAgainstSchema(inputSchema(), map[string]interface{}{ + "message": "hello", + "count": float64(3), + "mode": "fast", + "history": []interface{}{map[string]interface{}{"role": "user"}}, + }) + assert.Empty(t, errs, "a conforming payload must validate: %v", errs) +} + +func TestSchemaValidation_RejectsMissingRequiredField(t *testing.T) { + errs := validateAgainstSchema(inputSchema(), map[string]interface{}{"count": float64(1)}) + require.NotEmpty(t, errs) + assert.Contains(t, errs[0], "message") + assert.Contains(t, errs[0], "required") +} + +func TestSchemaValidation_RejectsWrongTypes(t *testing.T) { + cases := []struct { + name string + payload map[string]interface{} + want string + }{ + {"string field given a number", map[string]interface{}{"message": float64(5)}, "must be a string"}, + {"integer field given a string", map[string]interface{}{"message": "ok", "count": "many"}, "must be a number"}, + {"integer field given a fraction", map[string]interface{}{"message": "ok", "count": 1.5}, "must be an integer"}, + {"array field given an object", map[string]interface{}{"message": "ok", "history": map[string]interface{}{}}, "must be an array"}, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + errs := validateAgainstSchema(inputSchema(), tc.payload) + require.NotEmpty(t, errs) + assert.Contains(t, summariseErrors(errs), tc.want) + }) + } +} + +func TestSchemaValidation_EnforcesBounds(t *testing.T) { + tooShort := validateAgainstSchema(inputSchema(), map[string]interface{}{"message": ""}) + assert.Contains(t, summariseErrors(tooShort), "at least 1 characters") + + tooLong := validateAgainstSchema(inputSchema(), map[string]interface{}{"message": "far too long a message"}) + assert.Contains(t, summariseErrors(tooLong), "at most 10 characters") + + tooBig := validateAgainstSchema(inputSchema(), map[string]interface{}{"message": "ok", "count": float64(99)}) + assert.Contains(t, summariseErrors(tooBig), "at most") + + tooSmall := validateAgainstSchema(inputSchema(), map[string]interface{}{"message": "ok", "count": float64(-1)}) + assert.Contains(t, summariseErrors(tooSmall), "at least") +} + +func TestSchemaValidation_EnforcesEnum(t *testing.T) { + errs := validateAgainstSchema(inputSchema(), map[string]interface{}{"message": "ok", "mode": "sideways"}) + require.NotEmpty(t, errs) + assert.Contains(t, summariseErrors(errs), "must be one of") +} + +func TestSchemaValidation_ValidatesNestedItems(t *testing.T) { + errs := validateAgainstSchema(inputSchema(), map[string]interface{}{ + "message": "ok", + "history": []interface{}{ + map[string]interface{}{"role": "user"}, + map[string]interface{}{"content": "missing role"}, + }, + }) + require.NotEmpty(t, errs) + assert.Contains(t, summariseErrors(errs), "role") +} + +// A rule this validator does not implement must not make a valid payload fail. +func TestSchemaValidation_IgnoresUnsupportedKeywords(t *testing.T) { + schema := map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "email": map[string]interface{}{"type": "string", "format": "email", "pattern": "^.+@.+$"}, + }, + } + errs := validateAgainstSchema(schema, map[string]interface{}{"email": "not-an-email"}) + assert.Empty(t, errs, "unsupported keywords must be ignored, not reported as failures") +} + +func TestSchemaValidation_EmptySchemaAcceptsAnything(t *testing.T) { + assert.Empty(t, validateAgainstSchema(nil, map[string]interface{}{"anything": 1})) + assert.Empty(t, validateAgainstSchema(map[string]interface{}{}, "a string")) +} + +func TestSchemaValidation_RejectsNonObjectAtRoot(t *testing.T) { + errs := validateAgainstSchema(inputSchema(), "a bare string") + require.NotEmpty(t, errs) + assert.Contains(t, summariseErrors(errs), "must be an object") +} + +func TestSchemaValidation_AdditionalProperties(t *testing.T) { + schema := map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{"known": map[string]interface{}{"type": "string"}}, + "additionalProperties": false, + } + errs := validateAgainstSchema(schema, map[string]interface{}{"known": "x", "surprise": 1}) + require.NotEmpty(t, errs) + assert.Contains(t, summariseErrors(errs), "surprise") +} + +// Results must not depend on Go's map iteration order. +func TestSchemaValidation_IsDeterministic(t *testing.T) { + payload := map[string]interface{}{"count": "wrong", "mode": "sideways"} + + first := summariseErrors(validateAgainstSchema(inputSchema(), payload)) + for i := 0; i < 20; i++ { + assert.Equal(t, first, summariseErrors(validateAgainstSchema(inputSchema(), payload))) + } +} + +// Arbitrary values must never panic the validator. +func TestSchemaValidation_DoesNotPanic(t *testing.T) { + values := []interface{}{ + nil, "", 0, false, []interface{}{}, map[string]interface{}{}, + map[string]interface{}{"message": nil}, + map[string]interface{}{"history": []interface{}{nil, 1, "x"}}, + []interface{}{map[string]interface{}{"deep": []interface{}{map[string]interface{}{}}}}, + } + + for _, v := range values { + func() { + defer func() { + if r := recover(); r != nil { + t.Fatalf("validator panicked on %#v: %v", v, r) + } + }() + _ = validateAgainstSchema(inputSchema(), v) + }() + } +} diff --git a/pkg/server/security.go b/pkg/server/security.go new file mode 100644 index 0000000..0f05e1b --- /dev/null +++ b/pkg/server/security.go @@ -0,0 +1,210 @@ +// Copyright (c) 2024 GoLangGraph Team +// +// Licensed under the MIT License. See LICENSE file in the project root for full license information. +// +// Package: GoLangGraph - A powerful Go framework for building AI agent workflows + +package server + +import ( + "crypto/subtle" + "net/http" + "runtime/debug" + "strings" + "sync" + + "github.com/sirupsen/logrus" +) + +// SecurityConfig controls authentication, cross-origin access and request +// limits. The zero value is permissive so existing embedded uses keep working; +// production deployments should set RequireAuth and AllowedOrigins. +type SecurityConfig struct { + // RequireAuth rejects requests without a valid X-API-Key. + RequireAuth bool `json:"require_auth" yaml:"require_auth"` + // APIKeys are the accepted values for the X-API-Key header. + APIKeys []string `json:"api_keys" yaml:"api_keys"` + // AllowedOrigins restricts CORS and WebSocket origins. Empty means any + // origin is accepted, which is only appropriate for local development. + AllowedOrigins []string `json:"allowed_origins" yaml:"allowed_origins"` + // MaxRequestBytes caps request bodies. Zero applies DefaultMaxRequestBytes. + MaxRequestBytes int64 `json:"max_request_bytes" yaml:"max_request_bytes"` + // PublicPaths bypass authentication (health checks, readiness probes). + PublicPaths []string `json:"public_paths" yaml:"public_paths"` +} + +// DefaultMaxRequestBytes bounds request bodies so a single client cannot +// exhaust server memory with an oversized payload. +const DefaultMaxRequestBytes int64 = 4 << 20 // 4 MiB + +// DefaultSecurityConfig returns a development-friendly configuration: no auth, +// any origin, but with a request size limit already in place. +func DefaultSecurityConfig() *SecurityConfig { + return &SecurityConfig{ + RequireAuth: false, + MaxRequestBytes: DefaultMaxRequestBytes, + PublicPaths: []string{"/api/v1/health", "/health"}, + } +} + +// maxBytes returns the effective request size limit. +func (c *SecurityConfig) maxBytes() int64 { + if c == nil || c.MaxRequestBytes <= 0 { + return DefaultMaxRequestBytes + } + return c.MaxRequestBytes +} + +// isPublic reports whether a path bypasses authentication. +func (c *SecurityConfig) isPublic(path string) bool { + if c == nil { + return false + } + for _, p := range c.PublicPaths { + if p == path { + return true + } + } + return false +} + +// authorized reports whether a presented key matches a configured key. The +// comparison is constant time so a caller cannot recover a key by timing. +func (c *SecurityConfig) authorized(presented string) bool { + if c == nil || !c.RequireAuth { + return true + } + if presented == "" || len(c.APIKeys) == 0 { + return false + } + var ok bool + for _, key := range c.APIKeys { + if subtle.ConstantTimeCompare([]byte(key), []byte(presented)) == 1 { + ok = true + } + } + return ok +} + +// originAllowed reports whether an Origin header may access the API. +// +// An empty AllowedOrigins list accepts any origin, matching the previous +// permissive behavior for local development. A request with no Origin header +// is not a browser cross-origin request and is always allowed. +func (c *SecurityConfig) originAllowed(origin string) bool { + if origin == "" { + return true + } + if c == nil || len(c.AllowedOrigins) == 0 { + return true + } + for _, allowed := range c.AllowedOrigins { + if allowed == "*" { + return true + } + if strings.EqualFold(allowed, origin) { + return true + } + } + return false +} + +// corsOrigin returns the value to echo in Access-Control-Allow-Origin. +func (c *SecurityConfig) corsOrigin(origin string) string { + if c == nil || len(c.AllowedOrigins) == 0 { + return "*" + } + if c.originAllowed(origin) && origin != "" { + return origin + } + return "" +} + +// --------------------------------------------------------------------------- +// Middleware +// --------------------------------------------------------------------------- + +// recoveryMiddleware converts a panicking handler into a 500 response instead +// of tearing down the connection, and logs the stack for diagnosis. The stack +// is never written to the response. +func recoveryMiddleware(logger *logrus.Logger) func(http.Handler) http.Handler { + return func(next http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + defer func() { + if rec := recover(); rec != nil { + if logger != nil { + logger.WithFields(logrus.Fields{ + "panic": rec, + "path": r.URL.Path, + "method": r.Method, + "stack": string(debug.Stack()), + }).Error("Recovered from handler panic") + } + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusInternalServerError) + _, _ = w.Write([]byte(`{"error":"internal server error"}`)) + } + }() + next.ServeHTTP(w, r) + }) + } +} + +// bodyLimitMiddleware caps request body size. +func bodyLimitMiddleware(limit int64) func(http.Handler) http.Handler { + return func(next http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Body != nil { + r.Body = http.MaxBytesReader(w, r.Body, limit) + } + next.ServeHTTP(w, r) + }) + } +} + +// securityHeadersMiddleware sets conservative response headers. +func securityHeadersMiddleware(next http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("X-Content-Type-Options", "nosniff") + w.Header().Set("X-Frame-Options", "DENY") + w.Header().Set("Referrer-Policy", "no-referrer") + next.ServeHTTP(w, r) + }) +} + +// --------------------------------------------------------------------------- +// Concurrency-safe WebSocket writing +// --------------------------------------------------------------------------- + +// wsWriter serializes writes to a WebSocket connection. +// +// gorilla/websocket permits only one concurrent writer; the streaming handlers +// write from a goroutine while the read loop continues, so without this the +// connection can interleave frames and corrupt the stream. +type wsWriter struct { + mu sync.Mutex + conn wsConn +} + +// wsConn is the subset of *websocket.Conn used by the writer, which keeps this +// testable without a live connection. +type wsConn interface { + WriteJSON(v interface{}) error + Close() error +} + +func newWSWriter(conn wsConn) *wsWriter { return &wsWriter{conn: conn} } + +// WriteJSON writes a message, serialized against other writers. +func (w *wsWriter) WriteJSON(v interface{}) error { + w.mu.Lock() + defer w.mu.Unlock() + return w.conn.WriteJSON(v) +} + +// Close closes the underlying connection. +func (w *wsWriter) Close() error { + w.mu.Lock() + defer w.mu.Unlock() + return w.conn.Close() +} diff --git a/pkg/server/security_test.go b/pkg/server/security_test.go new file mode 100644 index 0000000..ba202fe --- /dev/null +++ b/pkg/server/security_test.go @@ -0,0 +1,618 @@ +// Copyright (c) 2024 GoLangGraph Team +// +// Licensed under the MIT License. See LICENSE file in the project root for full license information. + +package server + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "net/http" + "net/http/httptest" + "strings" + "sync" + "testing" + "time" + + "github.com/UnicoLab/GoLangGraph/pkg/core" + "github.com/UnicoLab/GoLangGraph/pkg/persistence" + "github.com/sirupsen/logrus" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// newTestServer builds a server with a registered graph for API tests. +func newTestServer(t *testing.T, mutate func(*ServerConfig)) *Server { + t.Helper() + cfg := DefaultServerConfig() + cfg.Port = 0 + // No static catch-all: it would shadow routes registered by individual tests. + cfg.StaticDir = "" + if mutate != nil { + mutate(cfg) + } + // The logger level now comes from cfg.LogLevel; overriding it here would + // mask whether that configuration is actually applied. + s := NewServer(cfg) + + g := core.NewGraph("demo") + g.Config.EnableStreaming = false + g.AddNode("greet", "Greet", func(ctx context.Context, st *core.BaseState) (*core.BaseState, error) { + in, _ := st.Get("input") + st.Set("greeting", fmt.Sprintf("hello %v", in)) + return st, nil + }) + g.AddNode("finish", "Finish", func(ctx context.Context, st *core.BaseState) (*core.BaseState, error) { + st.Set("done", true) + return st, nil + }) + g.AddEdge("greet", "finish", nil) + require.NoError(t, g.SetStartNode("greet")) + require.NoError(t, g.AddEndNode("finish")) + s.GraphManager().Register("demo", g) + + return s +} + +func doSecurityRequest(t *testing.T, s *Server, method, path string, body interface{}, headers map[string]string) *httptest.ResponseRecorder { + t.Helper() + var reader *bytes.Reader + if body != nil { + raw, err := json.Marshal(body) + require.NoError(t, err) + reader = bytes.NewReader(raw) + } else { + reader = bytes.NewReader(nil) + } + req := httptest.NewRequest(method, path, reader) + for k, v := range headers { + req.Header.Set(k, v) + } + rec := httptest.NewRecorder() + s.router.ServeHTTP(rec, req) + return rec +} + +// --------------------------------------------------------------------------- +// Authentication +// --------------------------------------------------------------------------- + +func TestServer_AuthRequiredRejectsMissingKey(t *testing.T) { + s := newTestServer(t, func(c *ServerConfig) { + c.Security.RequireAuth = true + c.Security.APIKeys = []string{"secret-key"} + }) + + rec := doSecurityRequest(t, s, http.MethodGet, "/api/v1/graphs", nil, nil) + assert.Equal(t, http.StatusUnauthorized, rec.Code) + assert.Contains(t, rec.Body.String(), "error") + + rec = doSecurityRequest(t, s, http.MethodGet, "/api/v1/graphs", nil, map[string]string{"X-API-Key": "wrong"}) + assert.Equal(t, http.StatusUnauthorized, rec.Code) + + rec = doSecurityRequest(t, s, http.MethodGet, "/api/v1/graphs", nil, map[string]string{"X-API-Key": "secret-key"}) + assert.Equal(t, http.StatusOK, rec.Code) +} + +func TestServer_HealthIsPublicUnderAuth(t *testing.T) { + s := newTestServer(t, func(c *ServerConfig) { + c.Security.RequireAuth = true + c.Security.APIKeys = []string{"k"} + }) + rec := doSecurityRequest(t, s, http.MethodGet, "/api/v1/health", nil, nil) + assert.Equal(t, http.StatusOK, rec.Code, "probes must not need credentials") +} + +func TestServer_AuthDisabledByDefault(t *testing.T) { + s := newTestServer(t, nil) + rec := doSecurityRequest(t, s, http.MethodGet, "/api/v1/graphs", nil, nil) + assert.Equal(t, http.StatusOK, rec.Code) +} + +func TestServer_AuthAcceptsAnyConfiguredKey(t *testing.T) { + s := newTestServer(t, func(c *ServerConfig) { + c.Security.RequireAuth = true + c.Security.APIKeys = []string{"first", "second"} + }) + for _, key := range []string{"first", "second"} { + rec := doSecurityRequest(t, s, http.MethodGet, "/api/v1/graphs", nil, map[string]string{"X-API-Key": key}) + assert.Equal(t, http.StatusOK, rec.Code, "key %q must be accepted", key) + } +} + +// An empty key list with auth on must fail closed rather than allowing all. +func TestServer_AuthWithNoKeysFailsClosed(t *testing.T) { + s := newTestServer(t, func(c *ServerConfig) { + c.Security.RequireAuth = true + c.Security.APIKeys = nil + }) + rec := doSecurityRequest(t, s, http.MethodGet, "/api/v1/graphs", nil, map[string]string{"X-API-Key": "anything"}) + assert.Equal(t, http.StatusUnauthorized, rec.Code) +} + +// --------------------------------------------------------------------------- +// CORS and origin handling +// --------------------------------------------------------------------------- + +func TestServer_CORSRestrictsOrigins(t *testing.T) { + s := newTestServer(t, func(c *ServerConfig) { + c.Security.AllowedOrigins = []string{"https://studio.example.com"} + }) + + rec := doSecurityRequest(t, s, http.MethodGet, "/api/v1/health", nil, + map[string]string{"Origin": "https://studio.example.com"}) + assert.Equal(t, "https://studio.example.com", rec.Header().Get("Access-Control-Allow-Origin")) + assert.Contains(t, rec.Header().Get("Vary"), "Origin") + + rec = doSecurityRequest(t, s, http.MethodGet, "/api/v1/health", nil, + map[string]string{"Origin": "https://evil.example.com"}) + assert.Empty(t, rec.Header().Get("Access-Control-Allow-Origin"), + "a disallowed origin must not receive CORS approval") +} + +func TestServer_CORSPreflightRejectsDisallowedOrigin(t *testing.T) { + s := newTestServer(t, func(c *ServerConfig) { + c.Security.AllowedOrigins = []string{"https://ok.example.com"} + }) + rec := doSecurityRequest(t, s, http.MethodOptions, "/api/v1/graphs", nil, + map[string]string{"Origin": "https://evil.example.com"}) + assert.Equal(t, http.StatusForbidden, rec.Code) +} + +func TestServer_CORSAllowsAPIKeyHeader(t *testing.T) { + s := newTestServer(t, nil) + rec := doSecurityRequest(t, s, http.MethodOptions, "/api/v1/graphs", nil, + map[string]string{"Origin": "http://localhost:5173"}) + assert.Equal(t, http.StatusOK, rec.Code) + assert.Contains(t, rec.Header().Get("Access-Control-Allow-Headers"), "X-API-Key", + "Studio authenticates with X-API-Key, so preflight must permit it") +} + +func TestServer_WebSocketOriginCheck(t *testing.T) { + s := newTestServer(t, func(c *ServerConfig) { + c.Security.AllowedOrigins = []string{"https://studio.example.com"} + }) + + allowed := httptest.NewRequest(http.MethodGet, "/api/v1/ws/graphs/demo/stream", nil) + allowed.Header.Set("Origin", "https://studio.example.com") + assert.True(t, s.upgrader.CheckOrigin(allowed)) + + denied := httptest.NewRequest(http.MethodGet, "/api/v1/ws/graphs/demo/stream", nil) + denied.Header.Set("Origin", "https://evil.example.com") + assert.False(t, s.upgrader.CheckOrigin(denied), + "accepting any origin allows cross-site WebSocket hijacking") +} + +// --------------------------------------------------------------------------- +// Request hygiene +// --------------------------------------------------------------------------- + +func TestServer_RejectsOversizedBody(t *testing.T) { + s := newTestServer(t, func(c *ServerConfig) { + c.Security.MaxRequestBytes = 1024 + }) + + huge := strings.Repeat("a", 8192) + req := httptest.NewRequest(http.MethodPost, "/api/v1/graphs/demo/execute", + strings.NewReader(`{"input":"`+huge+`"}`)) + req.Header.Set("Content-Type", "application/json") + rec := httptest.NewRecorder() + s.router.ServeHTTP(rec, req) + + assert.NotEqual(t, http.StatusOK, rec.Code, "an oversized body must be rejected") +} + +func TestServer_MalformedJSONIsRejected(t *testing.T) { + s := newTestServer(t, nil) + for _, body := range []string{"{", "", "null", "[]", `{"input":`, `{"input": 12345}`} { + req := httptest.NewRequest(http.MethodPost, "/api/v1/graphs/demo/execute", strings.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + rec := httptest.NewRecorder() + s.router.ServeHTTP(rec, req) + assert.NotEqual(t, http.StatusInternalServerError, rec.Code, + "malformed body %q must not produce a server error", body) + } +} + +func TestServer_SecurityHeadersArePresent(t *testing.T) { + s := newTestServer(t, nil) + rec := doSecurityRequest(t, s, http.MethodGet, "/api/v1/health", nil, nil) + assert.Equal(t, "nosniff", rec.Header().Get("X-Content-Type-Options")) + assert.Equal(t, "DENY", rec.Header().Get("X-Frame-Options")) +} + +func TestServer_PanicInHandlerReturns500(t *testing.T) { + s := newTestServer(t, nil) + s.router.HandleFunc("/boom", func(w http.ResponseWriter, r *http.Request) { + panic("handler exploded") + }) + + rec := doSecurityRequest(t, s, http.MethodGet, "/boom", nil, nil) + assert.Equal(t, http.StatusInternalServerError, rec.Code) + assert.NotContains(t, rec.Body.String(), "goroutine", "stack traces must not reach clients") + assert.Contains(t, rec.Body.String(), "error") +} + +// --------------------------------------------------------------------------- +// Graph API +// --------------------------------------------------------------------------- + +func TestServer_GraphListAndTopology(t *testing.T) { + s := newTestServer(t, nil) + + rec := doSecurityRequest(t, s, http.MethodGet, "/api/v1/graphs", nil, nil) + require.Equal(t, http.StatusOK, rec.Code) + var listed struct { + Graphs []GraphSummaryView `json:"graphs"` + } + require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &listed)) + require.Len(t, listed.Graphs, 1) + assert.Equal(t, "demo", listed.Graphs[0].ID) + assert.Equal(t, 2, listed.Graphs[0].NodeCount) + + rec = doSecurityRequest(t, s, http.MethodGet, "/api/v1/graphs/demo/topology", nil, nil) + require.Equal(t, http.StatusOK, rec.Code) + var topo struct { + GraphID string `json:"graph_id"` + Topology GraphTopologyView `json:"topology"` + } + require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &topo)) + assert.Equal(t, "demo", topo.GraphID) + require.Len(t, topo.Topology.Nodes, 2, "topology must report real nodes, not a placeholder") + require.Len(t, topo.Topology.Edges, 1) + assert.Equal(t, "greet", topo.Topology.Edges[0].From) + assert.Equal(t, "finish", topo.Topology.Edges[0].To) + + // Nodes carry the flags a visualiser needs. + byID := map[string]GraphNodeView{} + for _, n := range topo.Topology.Nodes { + byID[n.ID] = n + } + assert.True(t, byID["greet"].IsStart) + assert.True(t, byID["finish"].IsEnd) +} + +func TestServer_GraphNotFound(t *testing.T) { + s := newTestServer(t, nil) + for _, path := range []string{ + "/api/v1/graphs/missing", + "/api/v1/graphs/missing/topology", + } { + rec := doSecurityRequest(t, s, http.MethodGet, path, nil, nil) + assert.Equal(t, http.StatusNotFound, rec.Code, "path %s", path) + } + rec := doSecurityRequest(t, s, http.MethodPost, "/api/v1/graphs/missing/execute", + map[string]string{"input": "x"}, nil) + assert.Equal(t, http.StatusNotFound, rec.Code) +} + +func TestServer_GraphExecuteReturnsRealResult(t *testing.T) { + s := newTestServer(t, nil) + + rec := doSecurityRequest(t, s, http.MethodPost, "/api/v1/graphs/demo/execute", + map[string]interface{}{"input": "world"}, nil) + require.Equal(t, http.StatusOK, rec.Code) + + var resp struct { + GraphID string `json:"graph_id"` + Status string `json:"status"` + State map[string]interface{} `json:"state"` + Steps []ExecutionStepView `json:"steps"` + } + require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &resp)) + + assert.Equal(t, "completed", resp.Status) + assert.Equal(t, "hello world", resp.State["greeting"], + "execution must actually run the graph, not return a placeholder") + assert.Equal(t, true, resp.State["done"]) + require.Len(t, resp.Steps, 2, "every executed node must be reported") + assert.Equal(t, "greet", resp.Steps[0].NodeID) + assert.Equal(t, "finish", resp.Steps[1].NodeID) +} + +func TestServer_GraphExecuteReportsFailure(t *testing.T) { + s := newTestServer(t, nil) + failing := core.NewGraph("failing") + failing.Config.EnableStreaming = false + failing.AddNode("bad", "Bad", func(ctx context.Context, st *core.BaseState) (*core.BaseState, error) { + return nil, fmt.Errorf("provider unavailable") + }) + require.NoError(t, failing.SetStartNode("bad")) + s.GraphManager().Register("failing", failing) + + rec := doSecurityRequest(t, s, http.MethodPost, "/api/v1/graphs/failing/execute", + map[string]interface{}{"input": "x"}, nil) + assert.Equal(t, http.StatusInternalServerError, rec.Code) + + var resp map[string]interface{} + require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &resp)) + assert.Equal(t, "failed", resp["status"]) + assert.Contains(t, resp["error"], "provider unavailable") + assert.NotContains(t, fmt.Sprint(resp["error"]), "goroutine") +} + +func TestServer_GraphExecuteReportsInterruptAsResumable(t *testing.T) { + s := newTestServer(t, nil) + g := core.NewGraph("pausing") + g.Config.EnableStreaming = false + g.Config.InterruptBefore = []string{"second"} + g.AddNode("first", "First", func(ctx context.Context, st *core.BaseState) (*core.BaseState, error) { + st.Set("first", true) + return st, nil + }) + g.AddNode("second", "Second", func(ctx context.Context, st *core.BaseState) (*core.BaseState, error) { + st.Set("second", true) + return st, nil + }) + g.AddEdge("first", "second", nil) + require.NoError(t, g.SetStartNode("first")) + require.NoError(t, g.AddEndNode("second")) + s.GraphManager().Register("pausing", g) + + rec := doSecurityRequest(t, s, http.MethodPost, "/api/v1/graphs/pausing/execute", + map[string]interface{}{"input": "x"}, nil) + require.Equal(t, http.StatusOK, rec.Code, "a pause is a normal outcome, not a server error") + + var resp map[string]interface{} + require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &resp)) + assert.Equal(t, "interrupted", resp["status"]) + interrupt := resp["interrupt"].(map[string]interface{}) + assert.Equal(t, "second", interrupt["node_id"]) + assert.Equal(t, true, interrupt["before"]) + state := resp["state"].(map[string]interface{}) + assert.Equal(t, true, state["first"]) + assert.NotContains(t, state, "second") +} + +func TestServer_GraphExecuteAcceptsStateSeed(t *testing.T) { + s := newTestServer(t, nil) + rec := doSecurityRequest(t, s, http.MethodPost, "/api/v1/graphs/demo/execute", + map[string]interface{}{"state": map[string]interface{}{"input": "seeded", "extra": 1}}, nil) + require.Equal(t, http.StatusOK, rec.Code) + + var resp map[string]interface{} + require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &resp)) + state := resp["state"].(map[string]interface{}) + assert.Equal(t, "hello seeded", state["greeting"]) + assert.EqualValues(t, 1, state["extra"], "caller-supplied keys must reach the graph") +} + +// Concurrent API executions must not interfere with each other. +func TestServer_ConcurrentGraphExecutions(t *testing.T) { + s := newTestServer(t, nil) + + const n = 24 + var wg sync.WaitGroup + results := make([]string, n) + for i := 0; i < n; i++ { + wg.Add(1) + go func(i int) { + defer wg.Done() + rec := doSecurityRequest(t, s, http.MethodPost, "/api/v1/graphs/demo/execute", + map[string]interface{}{"input": fmt.Sprintf("client-%d", i)}, nil) + if rec.Code != http.StatusOK { + results[i] = "status " + rec.Result().Status + return + } + var resp map[string]interface{} + if err := json.Unmarshal(rec.Body.Bytes(), &resp); err != nil { + results[i] = err.Error() + return + } + state := resp["state"].(map[string]interface{}) + want := fmt.Sprintf("hello client-%d", i) + if state["greeting"] != want { + results[i] = fmt.Sprintf("got %v want %v", state["greeting"], want) + } + }(i) + } + wg.Wait() + + for i, r := range results { + assert.Empty(t, r, "request %d", i) + } +} + +// A canceled request must stop the run rather than finishing it. +func TestServer_RequestCancellationStopsExecution(t *testing.T) { + s := newTestServer(t, nil) + started := make(chan struct{}) + var once sync.Once + slow := core.NewGraph("slow") + slow.Config.EnableStreaming = false + slow.AddNode("wait", "Wait", func(ctx context.Context, st *core.BaseState) (*core.BaseState, error) { + once.Do(func() { close(started) }) + <-ctx.Done() + return nil, ctx.Err() + }) + require.NoError(t, slow.SetStartNode("wait")) + s.GraphManager().Register("slow", slow) + + ctx, cancel := context.WithCancel(context.Background()) + req := httptest.NewRequest(http.MethodPost, "/api/v1/graphs/slow/execute", + strings.NewReader(`{"input":"x"}`)).WithContext(ctx) + rec := httptest.NewRecorder() + + done := make(chan struct{}) + go func() { + defer close(done) + s.router.ServeHTTP(rec, req) + }() + + <-started + cancel() + + select { + case <-done: + case <-time.After(10 * time.Second): + t.Fatal("canceled request did not stop the graph run") + } + assert.Equal(t, http.StatusRequestTimeout, rec.Code) +} + +func TestServer_StopIsSafeBeforeStart(t *testing.T) { + s := newTestServer(t, nil) + assert.NoError(t, s.Stop(context.Background()), "Stop before Start must not panic") +} + +// --------------------------------------------------------------------------- +// Thread history and diagnostics +// --------------------------------------------------------------------------- + +// The checkpoints endpoint previously returned an empty list unconditionally, +// so a thread's history was invisible even when checkpoints existed. +func TestServer_ListCheckpointsReturnsRealHistory(t *testing.T) { + s := newTestServer(t, nil) + + cp := persistence.NewMemoryCheckpointer() + s.SetCheckpointer(cp) + + for step, node := range []string{"a", "b", "c"} { + st := core.NewBaseState() + st.Set("step", step) + require.NoError(t, cp.Save(context.Background(), &persistence.Checkpoint{ + ID: fmt.Sprintf("cp-%d", step), + ThreadID: "thread-1", + NodeID: node, + StepID: step, + State: st, + CreatedAt: time.Now().Add(time.Duration(step) * time.Second), + })) + } + + rec := doSecurityRequest(t, s, http.MethodGet, "/api/v1/threads/thread-1/checkpoints", nil, nil) + require.Equal(t, http.StatusOK, rec.Code) + + var resp struct { + ThreadID string `json:"thread_id"` + Checkpoints []struct { + ID string `json:"id"` + NodeID string `json:"node_id"` + StepID int `json:"step_id"` + } `json:"checkpoints"` + } + require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &resp)) + + assert.Equal(t, "thread-1", resp.ThreadID) + require.Len(t, resp.Checkpoints, 3, "the thread's checkpoints must be listed") + assert.Equal(t, []string{"a", "b", "c"}, + []string{resp.Checkpoints[0].NodeID, resp.Checkpoints[1].NodeID, resp.Checkpoints[2].NodeID}, + "checkpoints must be ordered oldest first so a client can replay them") +} + +func TestServer_ListCheckpointsWithoutCheckpointer(t *testing.T) { + s := newTestServer(t, nil) + rec := doSecurityRequest(t, s, http.MethodGet, "/api/v1/threads/t/checkpoints", nil, nil) + assert.Equal(t, http.StatusServiceUnavailable, rec.Code, + "a missing checkpointer must be reported, not reported as an empty history") +} + +func TestServer_ListCheckpointsRejectsUnsafeThreadID(t *testing.T) { + s := newTestServer(t, nil) + s.SetCheckpointer(persistence.NewFileCheckpointer(t.TempDir())) + + rec := doSecurityRequest(t, s, http.MethodGet, "/api/v1/threads/..%2F..%2Fetc/checkpoints", nil, nil) + assert.NotEqual(t, http.StatusOK, rec.Code) +} + +// Metrics must be measured, not hardcoded. +func TestServer_DebugMetricsAreMeasured(t *testing.T) { + s := newTestServer(t, func(c *ServerConfig) { c.DevMode = true }) + + // Generate some traffic, including a failure. + for i := 0; i < 3; i++ { + doSecurityRequest(t, s, http.MethodGet, "/api/v1/health", nil, nil) + } + doSecurityRequest(t, s, http.MethodGet, "/api/v1/graphs/missing", nil, nil) + + rec := doSecurityRequest(t, s, http.MethodGet, "/debug/metrics", nil, nil) + require.Equal(t, http.StatusOK, rec.Code) + + var resp struct { + Metrics map[string]interface{} `json:"metrics"` + } + require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &resp)) + + total, ok := resp.Metrics["requests_total"].(float64) + require.True(t, ok) + assert.Greater(t, total, float64(3), "requests must actually be counted") + + failed, ok := resp.Metrics["requests_failed"].(float64) + require.True(t, ok) + assert.GreaterOrEqual(t, failed, float64(1), "failures must be counted") + + alloc, ok := resp.Metrics["memory_alloc_bytes"].(float64) + require.True(t, ok, "memory must be measured, not reported as N/A") + assert.Greater(t, alloc, float64(0)) + + assert.Contains(t, resp.Metrics, "goroutines") + assert.Contains(t, resp.Metrics, "uptime_seconds") + assert.EqualValues(t, 1, resp.Metrics["graphs_registered"]) +} + +// Metrics must not panic when no agent manager is configured. +func TestServer_DebugMetricsWithoutAgentManager(t *testing.T) { + s := newTestServer(t, func(c *ServerConfig) { c.DevMode = true }) + s.SetAgentManager(nil) + + rec := doSecurityRequest(t, s, http.MethodGet, "/debug/metrics", nil, nil) + assert.Equal(t, http.StatusOK, rec.Code) +} + +// An endpoint that cannot do what it claims must say so rather than report +// success. +func TestServer_UnimplementedDebugEndpointsAreHonest(t *testing.T) { + s := newTestServer(t, func(c *ServerConfig) { c.DevMode = true }) + + reload := doSecurityRequest(t, s, http.MethodPost, "/debug/reload", nil, nil) + assert.Equal(t, http.StatusNotImplemented, reload.Code) + assert.NotContains(t, reload.Body.String(), "successfully", + "an operator must not be told a reload happened when none did") + + logs := doSecurityRequest(t, s, http.MethodGet, "/debug/logs", nil, nil) + assert.Equal(t, http.StatusNotImplemented, logs.Code) +} + +// Concurrent traffic must not race the metrics counters or the connection map. +func TestServer_MetricsAreRaceFree(t *testing.T) { + s := newTestServer(t, func(c *ServerConfig) { c.DevMode = true }) + + var wg sync.WaitGroup + for i := 0; i < 16; i++ { + wg.Add(1) + go func() { + defer wg.Done() + for j := 0; j < 10; j++ { + doSecurityRequest(t, s, http.MethodGet, "/api/v1/health", nil, nil) + doSecurityRequest(t, s, http.MethodGet, "/debug/metrics", nil, nil) + } + }() + } + wg.Wait() +} + +// LogLevel was declared in ServerConfig and never read, so an operator setting +// it saw no change in output at all. +func TestServer_LogLevelIsApplied(t *testing.T) { + for _, tc := range []struct { + configured string + want logrus.Level + }{ + {"debug", logrus.DebugLevel}, + {"warn", logrus.WarnLevel}, + {"error", logrus.ErrorLevel}, + } { + t.Run(tc.configured, func(t *testing.T) { + s := newTestServer(t, func(c *ServerConfig) { c.LogLevel = tc.configured }) + assert.Equal(t, tc.want, s.logger.GetLevel()) + }) + } +} + +// An unrecognized level must not change the logger or crash the server. +func TestServer_InvalidLogLevelKeepsDefault(t *testing.T) { + s := newTestServer(t, func(c *ServerConfig) { c.LogLevel = "not-a-level" }) + assert.Equal(t, logrus.InfoLevel, s.logger.GetLevel()) +} diff --git a/pkg/server/server.go b/pkg/server/server.go index f45f415..043263f 100644 --- a/pkg/server/server.go +++ b/pkg/server/server.go @@ -7,11 +7,17 @@ package server import ( + "bufio" "context" "encoding/json" + "errors" "fmt" + "net" "net/http" + "runtime" + "sort" "sync" + "sync/atomic" "time" "github.com/google/uuid" @@ -20,6 +26,7 @@ import ( "github.com/sirupsen/logrus" "github.com/UnicoLab/GoLangGraph/pkg/agent" + "github.com/UnicoLab/GoLangGraph/pkg/core" "github.com/UnicoLab/GoLangGraph/pkg/llm" "github.com/UnicoLab/GoLangGraph/pkg/persistence" "github.com/UnicoLab/GoLangGraph/pkg/tools" @@ -36,6 +43,10 @@ type ServerConfig struct { StaticDir string `json:"static_dir"` DevMode bool `json:"dev_mode"` LogLevel string `json:"log_level"` + + // Security controls authentication, allowed origins and request limits. + // Nil falls back to DefaultSecurityConfig. + Security *SecurityConfig `json:"security"` } // DefaultServerConfig returns default server configuration @@ -50,6 +61,7 @@ func DefaultServerConfig() *ServerConfig { StaticDir: "./static", DevMode: false, LogLevel: "info", + Security: DefaultSecurityConfig(), } } @@ -67,8 +79,17 @@ type Server struct { agentManager *AgentManager sessionManager *persistence.SessionManager - // WebSocket connections - wsConnections map[string]*websocket.Conn + graphManager *GraphManager + checkpointer persistence.Checkpointer + + // Request accounting for the metrics endpoint. + requestsTotal atomic.Uint64 + requestsFailed atomic.Uint64 + startedAt time.Time + + // WebSocket connections, keyed by resource ID then by connection, so that + // several clients can observe the same agent or graph at once. + wsConnections map[string]map[*websocket.Conn]struct{} wsConnectionsMu sync.RWMutex } @@ -78,22 +99,91 @@ func NewServer(config *ServerConfig) *Server { config = DefaultServerConfig() } + if config.Security == nil { + config.Security = DefaultSecurityConfig() + } + server := &Server{ config: config, router: mux.NewRouter(), logger: logrus.New(), - wsConnections: make(map[string]*websocket.Conn), - upgrader: websocket.Upgrader{ - CheckOrigin: func(r *http.Request) bool { - return true // Allow all origins for development - }, + graphManager: NewGraphManager(), + wsConnections: make(map[string]map[*websocket.Conn]struct{}), + startedAt: time.Now(), + } + + // Reject WebSocket upgrades from origins the API does not allow. Accepting + // every origin permits cross-site WebSocket hijacking: any page a user + // visits could open a socket to this server and drive it as that user. + server.upgrader = websocket.Upgrader{ + CheckOrigin: func(r *http.Request) bool { + return config.Security.originAllowed(r.Header.Get("Origin")) }, } + // LogLevel was declared in the configuration and never read, so setting it + // had no effect at all. + if config.LogLevel != "" { + if level, err := logrus.ParseLevel(config.LogLevel); err == nil { + server.logger.SetLevel(level) + } else { + server.logger.WithField("log_level", config.LogLevel). + Warn("Unrecognized log level; keeping the default") + } + } + server.setupRoutes() return server } +// SetCheckpointer attaches the checkpointer used to serve thread history. +func (s *Server) SetCheckpointer(cp persistence.Checkpointer) { + s.checkpointer = cp +} + +// SetGraphManager replaces the graph manager. +func (s *Server) SetGraphManager(manager *GraphManager) { + s.graphManager = manager +} + +// GraphManager returns the server's graph manager, used to register graphs +// that should be listable, inspectable and executable over the API. +func (s *Server) GraphManager() *GraphManager { + return s.graphManager +} + +// registerWSConn records a live WebSocket connection for a resource. +func (s *Server) registerWSConn(id string, conn *websocket.Conn) { + s.wsConnectionsMu.Lock() + defer s.wsConnectionsMu.Unlock() + if s.wsConnections[id] == nil { + s.wsConnections[id] = make(map[*websocket.Conn]struct{}) + } + s.wsConnections[id][conn] = struct{}{} +} + +// unregisterWSConn removes a connection, dropping the resource entry when the +// last connection for it closes. +func (s *Server) unregisterWSConn(id string, conn *websocket.Conn) { + s.wsConnectionsMu.Lock() + defer s.wsConnectionsMu.Unlock() + conns := s.wsConnections[id] + if conns == nil { + return + } + delete(conns, conn) + if len(conns) == 0 { + delete(s.wsConnections, id) + } +} + +// wsConnectionCount reports how many live connections a resource has. +func (s *Server) wsConnectionCount(id string) int { + s.wsConnectionsMu.RLock() + defer s.wsConnectionsMu.RUnlock() + return len(s.wsConnections[id]) +} + // SetLLMManager sets the LLM provider manager func (s *Server) SetLLMManager(manager *llm.ProviderManager) { s.llmManager = manager @@ -116,12 +206,20 @@ func (s *Server) SetSessionManager(manager *persistence.SessionManager) { // setupRoutes sets up HTTP routes func (s *Server) setupRoutes() { - // Enable CORS if configured + // Middleware. gorilla/mux wraps from the last registered to the first, so + // the first registered is the outermost. Recovery goes first so a panic in + // any later middleware or handler still produces a response rather than + // dropping the connection, and authentication goes last so a rejected + // request still carries the CORS and security headers a browser needs to + // read the 401. + s.router.Use(recoveryMiddleware(s.logger)) + s.router.Use(bodyLimitMiddleware(s.config.Security.maxBytes())) + s.router.Use(securityHeadersMiddleware) + if s.config.EnableCORS { s.router.Use(s.corsMiddleware) } - // Middleware s.router.Use(s.loggingMiddleware) s.router.Use(s.authMiddleware) @@ -183,6 +281,14 @@ func (s *Server) setupRoutes() { playground.HandleFunc("/agents/{id}/test", s.handlePlaygroundAgentTest).Methods("POST") } + // Cross-origin preflight. Routes declare concrete methods, so an OPTIONS + // request would otherwise fall through to 404 and the browser would block + // every cross-origin call. This catch-all gives the CORS middleware a + // matched route to run on; it must be registered before the static handler. + s.router.PathPrefix("/").Methods(http.MethodOptions).HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusNoContent) + }) + // Static files for UI if s.config.StaticDir != "" { s.router.PathPrefix("/").Handler(http.FileServer(http.Dir(s.config.StaticDir))) @@ -204,12 +310,30 @@ func (s *Server) Start() error { "port": s.config.Port, }).Info("Starting GoLangGraph server") - return s.server.ListenAndServe() + if err := s.server.ListenAndServe(); err != nil && !errors.Is(err, http.ErrServerClosed) { + return err + } + return nil } // Stop stops the server func (s *Server) Stop(ctx context.Context) error { s.logger.Info("Stopping GoLangGraph server") + + // Close live WebSocket connections so Shutdown is not blocked by hijacked + // connections, which http.Server does not wait for or close itself. + s.wsConnectionsMu.Lock() + for id, conns := range s.wsConnections { + for conn := range conns { + _ = conn.Close() + } + delete(s.wsConnections, id) + } + s.wsConnectionsMu.Unlock() + + if s.server == nil { + return nil + } return s.server.Shutdown(ctx) } @@ -217,11 +341,33 @@ func (s *Server) Stop(ctx context.Context) error { func (s *Server) corsMiddleware(next http.Handler) http.Handler { return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - w.Header().Set("Access-Control-Allow-Origin", "*") + origin := r.Header.Get("Origin") + + // Responses vary by Origin, so caches must not share them across origins. + w.Header().Add("Vary", "Origin") + + if !s.config.Security.originAllowed(origin) { + // Omit the CORS headers entirely; the browser then blocks the read. + if r.Method == http.MethodOptions { + w.WriteHeader(http.StatusForbidden) + return + } + next.ServeHTTP(w, r) + return + } + + if allow := s.config.Security.corsOrigin(origin); allow != "" { + w.Header().Set("Access-Control-Allow-Origin", allow) + // Credentials are only meaningful for a specific origin, never "*". + if allow != "*" { + w.Header().Set("Access-Control-Allow-Credentials", "true") + } + } w.Header().Set("Access-Control-Allow-Methods", "GET, POST, PUT, DELETE, OPTIONS") w.Header().Set("Access-Control-Allow-Headers", "Content-Type, Authorization, X-API-Key") + w.Header().Set("Access-Control-Max-Age", "600") - if r.Method == "OPTIONS" { + if r.Method == http.MethodOptions { w.WriteHeader(http.StatusOK) return } @@ -230,11 +376,47 @@ func (s *Server) corsMiddleware(next http.Handler) http.Handler { }) } +// statusRecorder captures the response status so metrics can count failures. +type statusRecorder struct { + http.ResponseWriter + status int +} + +func (r *statusRecorder) WriteHeader(status int) { + r.status = status + r.ResponseWriter.WriteHeader(status) +} + +// Flush and Hijack are forwarded so streaming and WebSocket upgrades still work +// through the recorder. +func (r *statusRecorder) Flush() { + if f, ok := r.ResponseWriter.(http.Flusher); ok { + f.Flush() + } +} + +func (r *statusRecorder) Hijack() (net.Conn, *bufio.ReadWriter, error) { + h, ok := r.ResponseWriter.(http.Hijacker) + if !ok { + return nil, nil, fmt.Errorf("response writer does not support hijacking") + } + return h.Hijack() +} + func (s *Server) loggingMiddleware(next http.Handler) http.Handler { return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { start := time.Now() + + s.requestsTotal.Add(1) + recorder := &statusRecorder{ResponseWriter: w, status: http.StatusOK} + w = recorder + next.ServeHTTP(w, r) + if recorder.status >= 400 { + s.requestsFailed.Add(1) + } + s.logger.WithFields(logrus.Fields{ "method": r.Method, "path": r.URL.Path, @@ -244,13 +426,28 @@ func (s *Server) loggingMiddleware(next http.Handler) http.Handler { }) } +// authMiddleware enforces API key authentication when the server is configured +// to require it. Preflight requests and configured public paths (health checks) +// bypass the check so probes and browsers keep working. func (s *Server) authMiddleware(next http.Handler) http.Handler { return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - // Simple authentication - in production, implement proper JWT/OAuth - // For now, just check for API key in header - apiKey := r.Header.Get("X-API-Key") - if apiKey == "" { - // Allow requests without API key for development + sec := s.config.Security + if sec == nil || !sec.RequireAuth { + next.ServeHTTP(w, r) + return + } + if r.Method == http.MethodOptions || sec.isPublic(r.URL.Path) { + next.ServeHTTP(w, r) + return + } + + if !sec.authorized(r.Header.Get("X-API-Key")) { + s.logger.WithFields(logrus.Fields{ + "path": r.URL.Path, + "remote": r.RemoteAddr, + }).Warn("Rejected unauthenticated request") + s.writeError(w, http.StatusUnauthorized, "missing or invalid API key") + return } next.ServeHTTP(w, r) @@ -284,9 +481,28 @@ func (s *Server) handleListProviders(w http.ResponseWriter, r *http.Request) { return } - providers := s.llmManager.ListProviders() + // Describe each provider rather than returning bare names, so clients can + // show its type, endpoint and model without a second round trip. + names := s.llmManager.ListProviders() + sort.Strings(names) + + infos := make([]map[string]interface{}, 0, len(names)) + for _, name := range names { + info := map[string]interface{}{"name": name} + if provider, err := s.llmManager.GetProvider(name); err == nil && provider != nil { + for key, value := range provider.GetConfig() { + // Never expose credentials over the API. + if key == "api_key" || key == "apiKey" { + continue + } + info[key] = value + } + } + infos = append(infos, info) + } + s.writeJSON(w, http.StatusOK, map[string]interface{}{ - "providers": providers, + "providers": infos, }) } @@ -352,13 +568,17 @@ func (s *Server) handleListAgents(w http.ResponseWriter, r *http.Request) { return } + // Return full configurations rather than bare IDs: clients such as + // GoLangGraph Studio render an agent's name, type, model and provider from + // this list, and a list of strings leaves every field undefined. ids := s.agentManager.ListAgents() configs := make([]*agent.AgentConfig, 0, len(ids)) for _, id := range ids { - if agentInstance, exists := s.agentManager.GetAgent(id); exists { - configs = append(configs, agentInstance.GetConfig()) + if instance, ok := s.agentManager.GetAgent(id); ok { + configs = append(configs, instance.GetConfig()) } } + sort.Slice(configs, func(i, j int) bool { return configs[i].ID < configs[j].ID }) s.writeJSON(w, http.StatusOK, map[string]interface{}{ "agents": configs, @@ -515,9 +735,16 @@ func (s *Server) handleGetAgentHistory(w http.ResponseWriter, r *http.Request) { } func (s *Server) handleListGraphs(w http.ResponseWriter, r *http.Request) { - // For now, return empty list - would need graph manager + summaries := []GraphSummaryView{} + if s.graphManager != nil { + for _, id := range s.graphManager.List() { + if g, ok := s.graphManager.Get(id); ok { + summaries = append(summaries, summariseGraph(id, g)) + } + } + } s.writeJSON(w, http.StatusOK, map[string]interface{}{ - "graphs": []string{}, + "graphs": summaries, }) } @@ -525,11 +752,18 @@ func (s *Server) handleGetGraph(w http.ResponseWriter, r *http.Request) { vars := mux.Vars(r) graphID := vars["id"] - // Placeholder implementation + graph, exists := s.lookupGraph(graphID) + if !exists { + s.writeError(w, http.StatusNotFound, "graph not found") + return + } + + topology := describeTopology(graph) s.writeJSON(w, http.StatusOK, map[string]interface{}{ "graph_id": graphID, - "nodes": []string{}, - "edges": []string{}, + "graph": summariseGraph(graphID, graph), + "nodes": topology.Nodes, + "edges": topology.Edges, }) } @@ -537,13 +771,15 @@ func (s *Server) handleGetGraphTopology(w http.ResponseWriter, r *http.Request) vars := mux.Vars(r) graphID := vars["id"] - // Placeholder implementation + graph, exists := s.lookupGraph(graphID) + if !exists { + s.writeError(w, http.StatusNotFound, "graph not found") + return + } + s.writeJSON(w, http.StatusOK, map[string]interface{}{ "graph_id": graphID, - "topology": map[string]interface{}{ - "nodes": []string{}, - "edges": []string{}, - }, + "topology": describeTopology(graph), }) } @@ -551,28 +787,89 @@ func (s *Server) handleExecuteGraph(w http.ResponseWriter, r *http.Request) { vars := mux.Vars(r) graphID := vars["id"] + graph, exists := s.lookupGraph(graphID) + if !exists { + s.writeError(w, http.StatusNotFound, "graph not found") + return + } + var request struct { - Input string `json:"input"` + Input string `json:"input"` + State map[string]core.StateValue `json:"state"` + ThreadID string `json:"thread_id"` } if err := json.NewDecoder(r.Body).Decode(&request); err != nil { - s.writeError(w, http.StatusBadRequest, "Invalid request body") + s.writeError(w, http.StatusBadRequest, "invalid request body: "+err.Error()) return } - // Placeholder implementation - s.writeJSON(w, http.StatusOK, map[string]interface{}{ - "graph_id": graphID, - "execution": "completed", - "result": "placeholder result", + steps := make(chan *core.ExecutionResult, 256) + collected := make([]ExecutionStepView, 0, 8) + drained := make(chan struct{}) + go func() { + defer close(drained) + for result := range steps { + collected = append(collected, describeStep(result)) + } + }() + + finalState, err := graph.ExecuteWithOptions(r.Context(), buildInitialState(request.Input, request.State), &core.ExecuteOptions{ + ThreadID: request.ThreadID, + Stream: steps, }) + <-drained + + response := map[string]interface{}{ + "graph_id": graphID, + "steps": collected, + } + if finalState != nil { + response["state"] = finalState.GetAll() + } + + if err != nil { + response["error"] = err.Error() + + var ie *core.InterruptError + switch { + case errors.As(err, &ie): + // A pause is a normal, resumable outcome, not a server error. + response["status"] = "interrupted" + response["interrupt"] = map[string]interface{}{ + "node_id": ie.NodeID, + "before": ie.Before, + "step": ie.Step, + } + s.writeJSON(w, http.StatusOK, response) + case errors.Is(err, context.Canceled), errors.Is(err, context.DeadlineExceeded): + response["status"] = "canceled" + s.writeJSON(w, http.StatusRequestTimeout, response) + case errors.Is(err, core.ErrGraphInvalid): + response["status"] = "invalid" + s.writeJSON(w, http.StatusBadRequest, response) + default: + response["status"] = "failed" + s.writeJSON(w, http.StatusInternalServerError, response) + } + return + } + + response["status"] = "completed" + s.writeJSON(w, http.StatusOK, response) } func (s *Server) handleInterruptGraph(w http.ResponseWriter, r *http.Request) { vars := mux.Vars(r) graphID := vars["id"] - // Placeholder implementation + graph, exists := s.lookupGraph(graphID) + if !exists { + s.writeError(w, http.StatusNotFound, "graph not found") + return + } + + graph.Interrupt() s.writeJSON(w, http.StatusOK, map[string]interface{}{ "graph_id": graphID, "status": "interrupted", @@ -702,10 +999,33 @@ func (s *Server) handleListCheckpoints(w http.ResponseWriter, r *http.Request) { vars := mux.Vars(r) threadID := vars["id"] - // Placeholder implementation + // Previously this always reported an empty list, so a thread's history was + // invisible even when checkpoints existed. + if s.checkpointer == nil { + s.writeError(w, http.StatusServiceUnavailable, "no checkpointer is configured") + return + } + + metadata, err := s.checkpointer.List(r.Context(), threadID) + if err != nil { + s.writeError(w, http.StatusBadRequest, err.Error()) + return + } + if metadata == nil { + metadata = []*persistence.CheckpointMetadata{} + } + + // Oldest first, so a client can replay a thread in order. + sort.Slice(metadata, func(i, j int) bool { + if metadata[i].StepID != metadata[j].StepID { + return metadata[i].StepID < metadata[j].StepID + } + return metadata[i].CreatedAt.Before(metadata[j].CreatedAt) + }) + s.writeJSON(w, http.StatusOK, map[string]interface{}{ "thread_id": threadID, - "checkpoints": []string{}, + "checkpoints": metadata, }) } @@ -741,6 +1061,12 @@ func (s *Server) handleGetTool(w http.ResponseWriter, r *http.Request) { }) } +// handleGraphWebSocket streams a graph execution to a client. +// +// Each connection gets its own execution context, canceled when the client +// disconnects, so a closed tab cannot leave a graph running forever. All writes +// go through a serializing writer because the read loop and the streaming +// goroutine write concurrently. func (s *Server) handleGraphWebSocket(w http.ResponseWriter, r *http.Request) { vars := mux.Vars(r) graphID := vars["id"] @@ -750,45 +1076,185 @@ func (s *Server) handleGraphWebSocket(w http.ResponseWriter, r *http.Request) { s.logger.WithError(err).Error("Failed to upgrade WebSocket") return } - defer conn.Close() + defer func() { _ = conn.Close() }() - // Store connection - s.wsConnectionsMu.Lock() - s.wsConnections[graphID] = conn - s.wsConnectionsMu.Unlock() + s.registerWSConn(graphID, conn) + defer s.unregisterWSConn(graphID, conn) - // Clean up on disconnect - defer func() { - s.wsConnectionsMu.Lock() - delete(s.wsConnections, graphID) - s.wsConnectionsMu.Unlock() - }() + writer := newWSWriter(conn) + + // Canceled when this connection goes away. + ctx, cancel := context.WithCancel(r.Context()) + defer cancel() + + var running sync.WaitGroup + defer running.Wait() - // Handle WebSocket messages for graph execution for { var message struct { - Type string `json:"type"` - Input string `json:"input"` + Type string `json:"type"` + Input string `json:"input"` + State map[string]core.StateValue `json:"state"` } - err := conn.ReadJSON(&message) - if err != nil { - s.logger.WithError(err).Error("WebSocket read error") + if err := conn.ReadJSON(&message); err != nil { + if !isFatalWSError(err) { + _ = writer.WriteJSON(map[string]interface{}{ + "type": "error", "graph_id": graphID, + "error": "invalid message: " + err.Error(), "timestamp": time.Now(), + }) + continue + } + if !websocket.IsCloseError(err, websocket.CloseNormalClosure, websocket.CloseGoingAway) { + s.logger.WithError(err).Debug("WebSocket read ended") + } + cancel() break } - // Placeholder graph execution - if message.Type == "execute" { - conn.WriteJSON(map[string]interface{}{ - "type": "result", - "graph_id": graphID, - "result": "Graph execution completed", + switch message.Type { + case "execute": + graph, exists := s.lookupGraph(graphID) + if !exists { + _ = writer.WriteJSON(map[string]interface{}{ + "type": "error", "graph_id": graphID, + "error": "graph not found", "timestamp": time.Now(), + }) + continue + } + running.Add(1) + go func() { + defer running.Done() + s.streamGraphExecution(ctx, writer, graphID, graph, message.Input, message.State) + }() + + case "interrupt": + if graph, exists := s.lookupGraph(graphID); exists { + graph.Interrupt() + _ = writer.WriteJSON(map[string]interface{}{ + "type": "interrupted", "graph_id": graphID, "timestamp": time.Now(), + }) + } + + case "ping": + _ = writer.WriteJSON(map[string]interface{}{"type": "pong", "timestamp": time.Now()}) + + default: + _ = writer.WriteJSON(map[string]interface{}{ + "type": "error", "error": "unknown message type: " + message.Type, "timestamp": time.Now(), }) } } } +// isFatalWSError reports whether a read error means the connection is gone, as +// opposed to a single malformed message. A client that sends one bad frame +// should get an error back and keep its session, not be disconnected. +func isFatalWSError(err error) bool { + if err == nil { + return false + } + var syntaxErr *json.SyntaxError + var typeErr *json.UnmarshalTypeError + if errors.As(err, &syntaxErr) || errors.As(err, &typeErr) { + // The frame was received in full; only its contents were unusable. + return false + } + return true +} + +// lookupGraph resolves a graph by ID. +// +// Registered graphs win, then the execution graph of an agent with that ID. +// Clients such as Studio request a topology using the agent's ID, so without +// the fallback the graph view of every agent would be empty. +func (s *Server) lookupGraph(id string) (*core.Graph, bool) { + if s.graphManager != nil { + if g, ok := s.graphManager.Get(id); ok { + return g, true + } + } + if s.agentManager != nil { + if instance, ok := s.agentManager.GetAgent(id); ok { + if g := instance.GetGraph(); g != nil { + return g, true + } + } + } + return nil, false +} + +// buildInitialState turns a WebSocket or HTTP request payload into a state. +func buildInitialState(input string, data map[string]core.StateValue) *core.BaseState { + state := core.NewBaseState() + for k, v := range data { + state.Set(k, v) + } + if input != "" { + state.Set("input", input) + } + return state +} + +// streamGraphExecution runs a graph and emits one message per node, then a +// terminal message. It never writes to the socket directly. +func (s *Server) streamGraphExecution(ctx context.Context, writer *wsWriter, graphID string, graph *core.Graph, input string, data map[string]core.StateValue) { + _ = writer.WriteJSON(map[string]interface{}{ + "type": "start", "graph_id": graphID, "timestamp": time.Now(), + }) + + steps := make(chan *core.ExecutionResult, 64) + done := make(chan struct{}) + + go func() { + defer close(done) + for result := range steps { + if err := writer.WriteJSON(map[string]interface{}{ + "type": "step", "graph_id": graphID, + "step": describeStep(result), "timestamp": time.Now(), + }); err != nil { + // The client is gone; drain the rest so the run is not blocked. + for range steps { + } + return + } + } + }() + + finalState, err := graph.ExecuteWithOptions(ctx, buildInitialState(input, data), &core.ExecuteOptions{ + Stream: steps, + }) + <-done + + if err != nil { + payload := map[string]interface{}{ + "type": "error", "graph_id": graphID, + "error": err.Error(), "timestamp": time.Now(), + } + var ie *core.InterruptError + if errors.As(err, &ie) { + payload["type"] = "interrupt" + payload["node_id"] = ie.NodeID + payload["before"] = ie.Before + payload["step"] = ie.Step + } + if finalState != nil { + payload["state"] = finalState.GetAll() + } + _ = writer.WriteJSON(payload) + return + } + + result := map[string]interface{}{ + "type": "complete", "graph_id": graphID, "timestamp": time.Now(), + } + if finalState != nil { + result["state"] = finalState.GetAll() + } + _ = writer.WriteJSON(result) +} + // WebSocket handlers func (s *Server) handleAgentWebSocket(w http.ResponseWriter, r *http.Request) { vars := mux.Vars(r) @@ -799,19 +1265,18 @@ func (s *Server) handleAgentWebSocket(w http.ResponseWriter, r *http.Request) { s.logger.WithError(err).Error("Failed to upgrade WebSocket") return } - defer conn.Close() + defer func() { _ = conn.Close() }() - // Store connection - s.wsConnectionsMu.Lock() - s.wsConnections[agentID] = conn - s.wsConnectionsMu.Unlock() + s.registerWSConn(agentID, conn) + defer s.unregisterWSConn(agentID, conn) - // Clean up on disconnect - defer func() { - s.wsConnectionsMu.Lock() - delete(s.wsConnections, agentID) - s.wsConnectionsMu.Unlock() - }() + writer := newWSWriter(conn) + + ctx, cancel := context.WithCancel(r.Context()) + defer cancel() + + var running sync.WaitGroup + defer running.Wait() // Handle WebSocket messages for { @@ -820,43 +1285,72 @@ func (s *Server) handleAgentWebSocket(w http.ResponseWriter, r *http.Request) { Input string `json:"input"` } - err := conn.ReadJSON(&message) - if err != nil { - s.logger.WithError(err).Error("WebSocket read error") + if err := conn.ReadJSON(&message); err != nil { + if !isFatalWSError(err) { + _ = writer.WriteJSON(map[string]interface{}{ + "type": "error", "error": "invalid message: " + err.Error(), "timestamp": time.Now(), + }) + continue + } + if !websocket.IsCloseError(err, websocket.CloseNormalClosure, websocket.CloseGoingAway) { + s.logger.WithError(err).Debug("WebSocket read ended") + } + cancel() break } - if message.Type == "execute" && s.agentManager != nil { + switch message.Type { + case "execute": + if s.agentManager == nil { + _ = writer.WriteJSON(map[string]interface{}{ + "type": "error", "error": "agent manager not available", "timestamp": time.Now(), + }) + continue + } agentInstance, exists := s.agentManager.GetAgent(agentID) - if exists { - go s.streamAgentExecution(conn, agentInstance, message.Input) + if !exists { + _ = writer.WriteJSON(map[string]interface{}{ + "type": "error", "error": "agent not found", "timestamp": time.Now(), + }) + continue } + running.Add(1) + go func(input string) { + defer running.Done() + s.streamAgentExecution(ctx, writer, agentInstance, input) + }(message.Input) + + case "ping": + _ = writer.WriteJSON(map[string]interface{}{"type": "pong", "timestamp": time.Now()}) + + default: + _ = writer.WriteJSON(map[string]interface{}{ + "type": "error", "error": "unknown message type: " + message.Type, "timestamp": time.Now(), + }) } } } -func (s *Server) streamAgentExecution(conn *websocket.Conn, agent agent.Agent, input string) { - ctx := context.Background() - - // Send start message - conn.WriteJSON(map[string]interface{}{ +// streamAgentExecution runs an agent for a WebSocket client. The context is the +// connection's, so a disconnect cancels the run instead of leaving it (and any +// provider calls it makes) running unattended. +func (s *Server) streamAgentExecution(ctx context.Context, writer *wsWriter, agentInstance agent.Agent, input string) { + _ = writer.WriteJSON(map[string]interface{}{ "type": "start", "timestamp": time.Now(), }) - // Execute agent - execution, err := agent.Execute(ctx, input) - + execution, err := agentInstance.Execute(ctx, input) if err != nil { - conn.WriteJSON(map[string]interface{}{ - "type": "error", - "error": err.Error(), + _ = writer.WriteJSON(map[string]interface{}{ + "type": "error", + "error": err.Error(), + "timestamp": time.Now(), }) return } - // Send result - conn.WriteJSON(map[string]interface{}{ + _ = writer.WriteJSON(map[string]interface{}{ "type": "result", "execution": execution, "timestamp": time.Now(), @@ -867,7 +1361,11 @@ func (s *Server) streamAgentExecution(conn *websocket.Conn, agent agent.Agent, i func (s *Server) writeJSON(w http.ResponseWriter, status int, data interface{}) { w.Header().Set("Content-Type", "application/json") w.WriteHeader(status) - json.NewEncoder(w).Encode(data) + if err := json.NewEncoder(w).Encode(data); err != nil { + // The status line is already written, so the response cannot be + // changed; record the failure instead of discarding it. + s.logger.WithError(err).Error("Failed to encode JSON response") + } } func (s *Server) writeError(w http.ResponseWriter, status int, message string) { @@ -951,34 +1449,59 @@ func (s *Server) handleDebugConfig(w http.ResponseWriter, r *http.Request) { } func (s *Server) handleDebugLogs(w http.ResponseWriter, r *http.Request) { - // In a real implementation, you would retrieve logs from a log store - s.writeJSON(w, http.StatusOK, map[string]interface{}{ - "logs": []map[string]interface{}{ - { - "timestamp": time.Now().Format(time.RFC3339), - "level": "info", - "message": "Debug logs endpoint accessed", - }, - }, + // No log store is wired up. Returning a fabricated entry would suggest logs + // are being served when none are; say so instead. + s.writeJSON(w, http.StatusNotImplemented, map[string]interface{}{ + "error": "no log store is configured; logs are written to the process logger", + "logs": []map[string]interface{}{}, }) } func (s *Server) handleDebugMetrics(w http.ResponseWriter, r *http.Request) { - // In a real implementation, you would collect actual metrics + // These were previously hardcoded, and the agent count dereferenced a nil + // manager while the WebSocket count read a map without its mutex. + agentsActive := 0 + if s.agentManager != nil { + agentsActive = len(s.agentManager.ListAgents()) + } + + s.wsConnectionsMu.RLock() + wsConnections := 0 + for _, conns := range s.wsConnections { + wsConnections += len(conns) + } + s.wsConnectionsMu.RUnlock() + + graphs := 0 + if s.graphManager != nil { + graphs = len(s.graphManager.List()) + } + + var mem runtime.MemStats + runtime.ReadMemStats(&mem) + s.writeJSON(w, http.StatusOK, map[string]interface{}{ "metrics": map[string]interface{}{ - "requests_total": 0, - "agents_active": len(s.agentManager.ListAgents()), - "websocket_connections": len(s.wsConnections), - "memory_usage": "N/A", + "requests_total": s.requestsTotal.Load(), + "requests_failed": s.requestsFailed.Load(), + "agents_active": agentsActive, + "graphs_registered": graphs, + "websocket_connections": wsConnections, + "goroutines": runtime.NumGoroutine(), + "memory_alloc_bytes": mem.Alloc, + "memory_sys_bytes": mem.Sys, + "gc_cycles": mem.NumGC, + "uptime_seconds": int64(time.Since(s.startedAt).Seconds()), }, }) } func (s *Server) handleDebugReload(w http.ResponseWriter, r *http.Request) { - // In a real implementation, you would reload configuration - s.writeJSON(w, http.StatusOK, map[string]interface{}{ - "message": "Configuration reloaded successfully", + // This reported "Configuration reloaded successfully" without reloading + // anything, which would lead an operator to believe a change had taken + // effect. Report the truth until reloading is actually implemented. + s.writeJSON(w, http.StatusNotImplemented, map[string]interface{}{ + "error": "configuration reload is not supported; restart the server to apply changes", "timestamp": time.Now().Format(time.RFC3339), }) } diff --git a/pkg/server/testdata/fuzz/FuzzAPIPaths/005945938d5a1787 b/pkg/server/testdata/fuzz/FuzzAPIPaths/005945938d5a1787 new file mode 100644 index 0000000..7d3ca5c --- /dev/null +++ b/pkg/server/testdata/fuzz/FuzzAPIPaths/005945938d5a1787 @@ -0,0 +1,2 @@ +go test fuzz v1 +string("00%") diff --git a/pkg/server/websocket_test.go b/pkg/server/websocket_test.go new file mode 100644 index 0000000..0524640 --- /dev/null +++ b/pkg/server/websocket_test.go @@ -0,0 +1,292 @@ +// Copyright (c) 2024 GoLangGraph Team +// +// Licensed under the MIT License. See LICENSE file in the project root for full license information. + +package server + +import ( + "context" + "fmt" + "net/http" + "net/http/httptest" + "strings" + "sync" + "testing" + "time" + + "github.com/UnicoLab/GoLangGraph/pkg/core" + "github.com/gorilla/websocket" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// wsTestServer starts a real HTTP server so WebSocket upgrades exercise the +// full path, not just the handler. +func wsTestServer(t *testing.T, mutate func(*ServerConfig)) (*Server, *httptest.Server) { + t.Helper() + s := newTestServer(t, mutate) + ts := httptest.NewServer(s.router) + t.Cleanup(ts.Close) + return s, ts +} + +func wsURL(ts *httptest.Server, path string) string { + return "ws" + strings.TrimPrefix(ts.URL, "http") + path +} + +func dialWS(t *testing.T, ts *httptest.Server, path string, headers http.Header) *websocket.Conn { + t.Helper() + conn, resp, err := websocket.DefaultDialer.Dial(wsURL(ts, path), headers) + if err != nil { + if resp != nil { + t.Fatalf("dial %s: %v (status %s)", path, err, resp.Status) + } + t.Fatalf("dial %s: %v", path, err) + } + t.Cleanup(func() { _ = conn.Close() }) + return conn +} + +// readMessages collects frames until a terminal type arrives or the deadline hits. +func readMessages(t *testing.T, conn *websocket.Conn, terminal map[string]bool, timeout time.Duration) []map[string]interface{} { + t.Helper() + require.NoError(t, conn.SetReadDeadline(time.Now().Add(timeout))) + + var messages []map[string]interface{} + for { + var msg map[string]interface{} + if err := conn.ReadJSON(&msg); err != nil { + return messages + } + messages = append(messages, msg) + if kind, _ := msg["type"].(string); terminal[kind] { + return messages + } + } +} + +func TestWebSocket_GraphExecutionStreamsEveryStep(t *testing.T) { + _, ts := wsTestServer(t, nil) + conn := dialWS(t, ts, "/api/v1/ws/graphs/demo/stream", nil) + + require.NoError(t, conn.WriteJSON(map[string]interface{}{ + "type": "execute", "input": "world", + })) + + messages := readMessages(t, conn, map[string]bool{"complete": true, "error": true}, 10*time.Second) + require.NotEmpty(t, messages) + + var kinds []string + var steps []string + var final map[string]interface{} + for _, m := range messages { + kind, _ := m["type"].(string) + kinds = append(kinds, kind) + switch kind { + case "step": + step := m["step"].(map[string]interface{}) + steps = append(steps, step["node_id"].(string)) + case "complete": + final, _ = m["state"].(map[string]interface{}) + } + } + + assert.Equal(t, "start", kinds[0]) + assert.Equal(t, "complete", kinds[len(kinds)-1], + "a run must finish with a terminal frame, got %v", kinds) + assert.Equal(t, []string{"greet", "finish"}, steps, + "every executed node must be streamed, in order") + require.NotNil(t, final) + assert.Equal(t, "hello world", final["greeting"]) +} + +func TestWebSocket_UnknownGraphReportsError(t *testing.T) { + _, ts := wsTestServer(t, nil) + conn := dialWS(t, ts, "/api/v1/ws/graphs/nope/stream", nil) + + require.NoError(t, conn.WriteJSON(map[string]interface{}{"type": "execute"})) + messages := readMessages(t, conn, map[string]bool{"error": true}, 5*time.Second) + require.NotEmpty(t, messages) + assert.Equal(t, "error", messages[len(messages)-1]["type"]) + assert.Contains(t, messages[len(messages)-1]["error"], "not found") +} + +func TestWebSocket_UnknownMessageTypeIsReported(t *testing.T) { + _, ts := wsTestServer(t, nil) + conn := dialWS(t, ts, "/api/v1/ws/graphs/demo/stream", nil) + + require.NoError(t, conn.WriteJSON(map[string]interface{}{"type": "not-a-command"})) + messages := readMessages(t, conn, map[string]bool{"error": true}, 5*time.Second) + require.NotEmpty(t, messages) + assert.Contains(t, messages[0]["error"], "unknown message type") +} + +func TestWebSocket_PingPong(t *testing.T) { + _, ts := wsTestServer(t, nil) + conn := dialWS(t, ts, "/api/v1/ws/graphs/demo/stream", nil) + + require.NoError(t, conn.WriteJSON(map[string]interface{}{"type": "ping"})) + messages := readMessages(t, conn, map[string]bool{"pong": true}, 5*time.Second) + require.NotEmpty(t, messages) + assert.Equal(t, "pong", messages[0]["type"]) +} + +// Malformed frames must not kill the server or leak goroutines. +func TestWebSocket_MalformedFramesAreSurvived(t *testing.T) { + _, ts := wsTestServer(t, nil) + conn := dialWS(t, ts, "/api/v1/ws/graphs/demo/stream", nil) + + // A single unusable frame must be answered with an error, not a disconnect. + require.NoError(t, conn.WriteMessage(websocket.TextMessage, []byte("{not json"))) + messages := readMessages(t, conn, map[string]bool{"error": true}, 5*time.Second) + require.NotEmpty(t, messages) + assert.Contains(t, messages[0]["error"], "invalid message") + + // A field of the wrong type is equally recoverable. + require.NoError(t, conn.WriteJSON(map[string]interface{}{"type": "execute", "input": 12345})) + messages = readMessages(t, conn, map[string]bool{"error": true}, 5*time.Second) + require.NotEmpty(t, messages) + + // The same session still works afterwards. + require.NoError(t, conn.WriteJSON(map[string]interface{}{"type": "ping"})) + messages = readMessages(t, conn, map[string]bool{"pong": true}, 5*time.Second) + require.NotEmpty(t, messages, "a malformed frame must not end the session") +} + +// The origin allowlist must apply to WebSocket upgrades. +func TestWebSocket_RejectsDisallowedOrigin(t *testing.T) { + _, ts := wsTestServer(t, func(c *ServerConfig) { + c.Security.AllowedOrigins = []string{"https://studio.example.com"} + }) + + headers := http.Header{} + headers.Set("Origin", "https://evil.example.com") + _, resp, err := websocket.DefaultDialer.Dial(wsURL(ts, "/api/v1/ws/graphs/demo/stream"), headers) + require.Error(t, err, "a disallowed origin must not complete the upgrade") + if resp != nil { + assert.Equal(t, http.StatusForbidden, resp.StatusCode) + } +} + +// Several clients may watch the same graph; one disconnecting must not evict +// the others from the connection registry. +func TestWebSocket_MultipleClientsPerGraph(t *testing.T) { + s, ts := wsTestServer(t, nil) + + first := dialWS(t, ts, "/api/v1/ws/graphs/demo/stream", nil) + second := dialWS(t, ts, "/api/v1/ws/graphs/demo/stream", nil) + + // Both must be registered. + require.Eventually(t, func() bool { return s.wsConnectionCount("demo") == 2 }, + 5*time.Second, 20*time.Millisecond, "both clients must be tracked") + + require.NoError(t, first.Close()) + require.Eventually(t, func() bool { return s.wsConnectionCount("demo") == 1 }, + 5*time.Second, 20*time.Millisecond, "closing one client must not drop the other") + + // The survivor still works. + require.NoError(t, second.WriteJSON(map[string]interface{}{"type": "ping"})) + messages := readMessages(t, second, map[string]bool{"pong": true}, 5*time.Second) + require.NotEmpty(t, messages) +} + +// Concurrent executions on one connection must not interleave frames. Without +// serialized writes gorilla/websocket corrupts the stream. +func TestWebSocket_ConcurrentExecutionsDoNotCorruptStream(t *testing.T) { + s, ts := wsTestServer(t, nil) + + // A graph slow enough that several runs overlap. + slow := core.NewGraph("slow-stream") + for i := 0; i < 6; i++ { + id := fmt.Sprintf("n%d", i) + slow.AddNode(id, id, func(ctx context.Context, st *core.BaseState) (*core.BaseState, error) { + time.Sleep(5 * time.Millisecond) + return st, nil + }) + if i > 0 { + slow.AddEdge(fmt.Sprintf("n%d", i-1), id, nil) + } + } + require.NoError(t, slow.SetStartNode("n0")) + require.NoError(t, slow.AddEndNode("n5")) + s.GraphManager().Register("slow-stream", slow) + + conn := dialWS(t, ts, "/api/v1/ws/graphs/slow-stream/stream", nil) + + const runs = 4 + for i := 0; i < runs; i++ { + require.NoError(t, conn.WriteJSON(map[string]interface{}{"type": "execute", "input": fmt.Sprintf("run-%d", i)})) + } + + require.NoError(t, conn.SetReadDeadline(time.Now().Add(30*time.Second))) + completes := 0 + frames := 0 + for completes < runs { + var msg map[string]interface{} + // A corrupted stream surfaces here as a JSON decode failure. + require.NoError(t, conn.ReadJSON(&msg), "stream corrupted after %d frames", frames) + frames++ + if kind, _ := msg["type"].(string); kind == "complete" { + completes++ + } + } + assert.Equal(t, runs, completes) +} + +// Closing the connection must cancel the run instead of leaving it going. +func TestWebSocket_DisconnectCancelsExecution(t *testing.T) { + s, ts := wsTestServer(t, nil) + + started := make(chan struct{}) + finished := make(chan error, 1) + var once sync.Once + + blocking := core.NewGraph("blocking") + blocking.AddNode("wait", "Wait", func(ctx context.Context, st *core.BaseState) (*core.BaseState, error) { + once.Do(func() { close(started) }) + <-ctx.Done() + finished <- ctx.Err() + return nil, ctx.Err() + }) + require.NoError(t, blocking.SetStartNode("wait")) + s.GraphManager().Register("blocking", blocking) + + conn := dialWS(t, ts, "/api/v1/ws/graphs/blocking/stream", nil) + require.NoError(t, conn.WriteJSON(map[string]interface{}{"type": "execute"})) + + select { + case <-started: + case <-time.After(10 * time.Second): + t.Fatal("graph never started") + } + + require.NoError(t, conn.Close()) + + select { + case err := <-finished: + assert.ErrorIs(t, err, context.Canceled, + "a disconnected client must cancel its run rather than leaving it running") + case <-time.After(15 * time.Second): + t.Fatal("run was not canceled when the client disconnected") + } +} + +// Server shutdown must close live WebSocket connections rather than hanging. +func TestWebSocket_ShutdownClosesConnections(t *testing.T) { + s := newTestServer(t, nil) + ts := httptest.NewServer(s.router) + defer ts.Close() + + conn := dialWS(t, ts, "/api/v1/ws/graphs/demo/stream", nil) + require.NoError(t, conn.WriteJSON(map[string]interface{}{"type": "ping"})) + require.NotEmpty(t, readMessages(t, conn, map[string]bool{"pong": true}, 5*time.Second)) + + require.Eventually(t, func() bool { return s.wsConnectionCount("demo") == 1 }, + 5*time.Second, 20*time.Millisecond) + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + require.NoError(t, s.Stop(ctx)) + + assert.Zero(t, s.wsConnectionCount("demo"), "shutdown must release tracked connections") +} diff --git a/pkg/tools/fuzz_test.go b/pkg/tools/fuzz_test.go new file mode 100644 index 0000000..ba16db3 --- /dev/null +++ b/pkg/tools/fuzz_test.go @@ -0,0 +1,183 @@ +// Copyright (c) 2024 GoLangGraph Team +// +// Licensed under the MIT License. See LICENSE file in the project root for full license information. + +package tools + +import ( + "context" + "path/filepath" + "strings" + "testing" +) + +// fuzzTools returns every built-in tool, confined to a sandbox so a fuzzed +// argument cannot touch anything outside the test's temporary directory. +func fuzzTools(t *testing.T, root string) []Tool { + policy := DefaultSecurityPolicy() + policy.SetAllowedRoots([]string{root}) + // Never let a fuzzed URL reach the network. + policy.AllowedHosts = []string{"127.0.0.1.invalid"} + + all := []Tool{ + NewFileReadTool(), NewFileWriteTool(), NewFileListTool(), + NewShellTool(), NewHTTPTool(), NewCalculatorTool(), + NewTimeTool(), NewWebSearchTool(), + } + for _, tool := range all { + if setter, ok := tool.(interface{ SetSecurityPolicy(*SecurityPolicy) }); ok { + setter.SetSecurityPolicy(policy) + } + } + return all +} + +// FuzzToolArguments feeds arbitrary payloads to every tool. Tool arguments are +// produced by a language model, so they are untrusted: a tool must return an +// error, never panic and never escape its sandbox. +func FuzzToolArguments(f *testing.F) { + seeds := []string{ + `{"file_path":"/etc/passwd"}`, + `{"file_path":"../../../../etc/shadow"}`, + `{"command":"ls -la"}`, + `{"url":"http://169.254.169.254/"}`, + `{"expression":"1+1"}`, + `{"query":"anything"}`, + `{"path":"."}`, + `{}`, `[]`, `null`, ``, `{`, + `{"file_path":" "}`, + `{"expression":"9999999999^9999999999"}`, + strings.Repeat(`{"a":`, 100), + } + for _, s := range seeds { + f.Add(s) + } + + f.Fuzz(func(t *testing.T, args string) { + root := t.TempDir() + resolved, err := filepath.EvalSymlinks(root) + if err != nil { + resolved = root + } + + for _, tool := range fuzzTools(t, resolved) { + func() { + defer func() { + if r := recover(); r != nil { + t.Fatalf("%s panicked on %q: %v", tool.GetName(), args, r) + } + }() + _ = tool.Validate(args) + // Any error is fine; a crash or an escape is not. + _, _ = tool.Execute(context.Background(), args) + }() + } + }) +} + +// FuzzPathResolution checks the filesystem confinement directly: no input may +// resolve to a path outside the allowed roots. +func FuzzPathResolution(f *testing.F) { + seeds := []string{ + "file.txt", "../escape", "../../../../etc/passwd", "/etc/passwd", + "./nested/../file", "", " ", "sub/dir/file.json", "/proc/self/environ", + strings.Repeat("../", 50) + "etc/passwd", + "embedded" + string(rune(0)) + "null", + } + for _, s := range seeds { + f.Add(s) + } + + f.Fuzz(func(t *testing.T, path string) { + root := t.TempDir() + resolved, err := filepath.EvalSymlinks(root) + if err != nil { + resolved = root + } + policy := DefaultSecurityPolicy() + policy.SetAllowedRoots([]string{resolved}) + + got, err := policy.ResolvePath(path) + if err != nil { + return // refusing the path is a correct outcome + } + + // An accepted path must lie within the root. + if got != resolved && !strings.HasPrefix(got, resolved+string(filepath.Separator)) { + t.Fatalf("path %q escaped the sandbox: resolved to %q, root %q", path, got, resolved) + } + if !filepath.IsAbs(got) { + t.Fatalf("resolved path %q is not absolute (input %q)", got, path) + } + }) +} + +// FuzzCommandPolicy checks that no argument string can smuggle execution past +// the shell allowlist. +func FuzzCommandPolicy(f *testing.F) { + seeds := []string{ + "echo hi", "ls -la", "rm -rf /", "/bin/sh -c id", + "find . -exec rm", "", " ", "echo a b c", + } + for _, s := range seeds { + f.Add(s) + } + + f.Fuzz(func(t *testing.T, command string) { + policy := DefaultSecurityPolicy() + parts := strings.Fields(command) + + if err := policy.CheckCommand(parts); err != nil { + return // refusing is correct + } + + // Anything accepted must be a bare allowlisted name. + if len(parts) == 0 { + t.Fatalf("an empty command was accepted") + } + base := parts[0] + permitted := false + for _, allowed := range policy.AllowedCommands { + if base == allowed { + permitted = true + break + } + } + if !permitted { + t.Fatalf("command %q was accepted but is not allowlisted", base) + } + if strings.ContainsAny(base, "/\\") { + t.Fatalf("command %q was accepted as a path", base) + } + }) +} + +// FuzzURLPolicy checks that no URL accepted by the policy points at a local or +// private address. +func FuzzURLPolicy(f *testing.F) { + seeds := []string{ + "http://example.com", "https://example.com/path", + "http://127.0.0.1", "http://169.254.169.254/latest/", + "file:///etc/passwd", "", "http://", + "http://2130706433/", "http://0x7f000001/", "http://localhost", + } + for _, s := range seeds { + f.Add(s) + } + + f.Fuzz(func(t *testing.T, raw string) { + policy := DefaultSecurityPolicy() + parsed, err := policy.CheckURL(raw) + if err != nil { + return // refusing is correct + } + + scheme := strings.ToLower(parsed.Scheme) + if scheme != "http" && scheme != "https" { + t.Fatalf("URL %q accepted with scheme %q", raw, parsed.Scheme) + } + if parsed.Host == "" { + t.Fatalf("URL %q accepted with no host", raw) + } + }) +} diff --git a/pkg/tools/security.go b/pkg/tools/security.go new file mode 100644 index 0000000..21ee704 --- /dev/null +++ b/pkg/tools/security.go @@ -0,0 +1,369 @@ +// Copyright (c) 2024 GoLangGraph Team +// +// Licensed under the MIT License. See LICENSE file in the project root for full license information. +// +// Package: GoLangGraph - A powerful Go framework for building AI agent workflows + +package tools + +import ( + "context" + "fmt" + "net" + "net/http" + "net/url" + "os" + "path/filepath" + "strings" + "sync" + "syscall" + "time" +) + +// SecurityPolicy bounds what the built-in tools may touch. +// +// Tool arguments are chosen by a language model, which may be steered by +// untrusted input, so every tool that reaches the filesystem, a shell or the +// network is treated as attacker-controlled and constrained here. +type SecurityPolicy struct { + mu sync.RWMutex + + // AllowedRoots are absolute directories the file tools may operate within. + // Empty means the process working directory and the system temp directory. + AllowedRoots []string + + // MaxOutputBytes caps command output and HTTP response bodies. + MaxOutputBytes int64 + + // AllowedCommands is the shell command allowlist. + AllowedCommands []string + + // DeniedArgSubstrings reject shell arguments that can execute other + // programs even when the base command is allowed. + DeniedArgSubstrings []string + + // AllowPrivateNetwork permits requests to loopback, link-local and private + // address ranges. Off by default: those ranges hold cloud metadata services + // and internal admin endpoints. + AllowPrivateNetwork bool + + // AllowedHosts, when non-empty, restricts HTTP requests to these hosts. + AllowedHosts []string + + // MaxRedirects bounds HTTP redirect following. + MaxRedirects int +} + +// DefaultMaxToolOutputBytes caps how much a single tool call may return. +const DefaultMaxToolOutputBytes int64 = 1 << 20 // 1 MiB + +// DefaultSecurityPolicy returns the policy applied to tools built with the +// New*Tool constructors. +func DefaultSecurityPolicy() *SecurityPolicy { + return &SecurityPolicy{ + MaxOutputBytes: DefaultMaxToolOutputBytes, + // "find" is deliberately absent: -exec, -execdir and -ok run arbitrary + // programs, which defeats any command allowlist. + AllowedCommands: []string{"ls", "pwd", "echo", "wc", "head", "tail"}, + DeniedArgSubstrings: []string{ + "-exec", "-execdir", "-ok", "-okdir", "-fprintf", "-delete", + "--eval", "-e", "$(", "`", ";", "|", "&&", "||", ">", "<", + }, + MaxRedirects: 5, + } +} + +// roots returns the effective allowed roots, resolved to absolute paths. +func (p *SecurityPolicy) roots() []string { + p.mu.RLock() + configured := append([]string(nil), p.AllowedRoots...) + p.mu.RUnlock() + + if len(configured) == 0 { + cwd, err := os.Getwd() + if err == nil { + configured = append(configured, cwd) + } + configured = append(configured, os.TempDir()) + } + + out := make([]string, 0, len(configured)) + for _, root := range configured { + abs, err := filepath.Abs(root) + if err != nil { + continue + } + if resolved, err := filepath.EvalSymlinks(abs); err == nil { + abs = resolved + } + out = append(out, abs) + } + return out +} + +// SetAllowedRoots replaces the directories file tools may operate within. +func (p *SecurityPolicy) SetAllowedRoots(roots []string) { + p.mu.Lock() + defer p.mu.Unlock() + p.AllowedRoots = append([]string(nil), roots...) +} + +// maxOutput returns the effective output cap. +func (p *SecurityPolicy) maxOutput() int64 { + if p == nil || p.MaxOutputBytes <= 0 { + return DefaultMaxToolOutputBytes + } + return p.MaxOutputBytes +} + +// ResolvePath validates a caller-supplied path and returns its cleaned absolute +// form. Symlinks are resolved so a link inside an allowed root cannot be used +// to reach a file outside one; for a path that does not exist yet, the nearest +// existing parent is checked instead so new files can still be created. +func (p *SecurityPolicy) ResolvePath(path string) (string, error) { + if strings.TrimSpace(path) == "" { + return "", fmt.Errorf("path must not be empty") + } + if strings.ContainsRune(path, 0) { + return "", fmt.Errorf("path must not contain a null byte") + } + + abs, err := filepath.Abs(path) + if err != nil { + return "", fmt.Errorf("invalid path: %w", err) + } + abs = filepath.Clean(abs) + + // Resolve symlinks on the deepest existing ancestor. + probe := abs + var trailing []string + for { + if resolved, err := filepath.EvalSymlinks(probe); err == nil { + probe = resolved + break + } + parent := filepath.Dir(probe) + if parent == probe { + break + } + trailing = append([]string{filepath.Base(probe)}, trailing...) + probe = parent + } + resolved := filepath.Join(append([]string{probe}, trailing...)...) + + roots := p.roots() + if len(roots) == 0 { + return "", fmt.Errorf("no allowed roots are configured") + } + for _, root := range roots { + if resolved == root || strings.HasPrefix(resolved, root+string(os.PathSeparator)) { + return resolved, nil + } + } + + return "", fmt.Errorf("path %q is outside the allowed directories", path) +} + +// CheckCommand validates a shell invocation against the allowlist and rejects +// arguments that would let an allowed command run something else. +func (p *SecurityPolicy) CheckCommand(parts []string) error { + if len(parts) == 0 { + return fmt.Errorf("empty command") + } + + p.mu.RLock() + allowed := append([]string(nil), p.AllowedCommands...) + denied := append([]string(nil), p.DeniedArgSubstrings...) + p.mu.RUnlock() + + base := parts[0] + if strings.ContainsAny(base, "/\\") { + return fmt.Errorf("command %q must be a bare name, not a path", base) + } + + permitted := false + for _, cmd := range allowed { + if base == cmd { + permitted = true + break + } + } + if !permitted { + return fmt.Errorf("command %s not allowed", base) + } + + for _, arg := range parts[1:] { + for _, bad := range denied { + if strings.Contains(arg, bad) { + return fmt.Errorf("argument %q is not allowed (contains %q)", arg, bad) + } + } + } + + return nil +} + +// --------------------------------------------------------------------------- +// Network policy +// --------------------------------------------------------------------------- + +// CheckURL validates a request target before any connection is made. +func (p *SecurityPolicy) CheckURL(raw string) (*url.URL, error) { + parsed, err := url.Parse(raw) + if err != nil { + return nil, fmt.Errorf("invalid URL: %w", err) + } + switch strings.ToLower(parsed.Scheme) { + case "http", "https": + default: + return nil, fmt.Errorf("URL scheme %q is not allowed", parsed.Scheme) + } + if parsed.Host == "" { + return nil, fmt.Errorf("URL must include a host") + } + + p.mu.RLock() + allowedHosts := append([]string(nil), p.AllowedHosts...) + allowPrivate := p.AllowPrivateNetwork + p.mu.RUnlock() + + host := parsed.Hostname() + + if len(allowedHosts) > 0 { + ok := false + for _, h := range allowedHosts { + if strings.EqualFold(h, host) { + ok = true + break + } + } + if !ok { + return nil, fmt.Errorf("host %q is not in the allowed list", host) + } + } + + // A literal IP can be checked now; hostnames are checked at dial time, + // which also defeats DNS rebinding. + if ip := net.ParseIP(host); ip != nil && !allowPrivate { + if err := checkPublicIP(ip); err != nil { + return nil, err + } + } + + return parsed, nil +} + +// checkPublicIP rejects addresses that reach the host itself or the local +// network, including cloud metadata endpoints. +func checkPublicIP(ip net.IP) error { + switch { + case ip.IsLoopback(): + return fmt.Errorf("address %s is loopback", ip) + case ip.IsPrivate(): + return fmt.Errorf("address %s is in a private range", ip) + case ip.IsLinkLocalUnicast(), ip.IsLinkLocalMulticast(): + // 169.254.169.254 is the cloud instance metadata service. + return fmt.Errorf("address %s is link-local", ip) + case ip.IsUnspecified(): + return fmt.Errorf("address %s is unspecified", ip) + case ip.IsMulticast(): + return fmt.Errorf("address %s is multicast", ip) + case ip.IsInterfaceLocalMulticast(): + return fmt.Errorf("address %s is interface-local", ip) + } + // Unique local IPv6 (fc00::/7) is not covered by IsPrivate on all versions. + if len(ip) == net.IPv6len && ip[0]&0xfe == 0xfc { + return fmt.Errorf("address %s is a unique local address", ip) + } + return nil +} + +// HTTPClient builds a client that enforces the policy on every connection and +// every redirect hop. +func (p *SecurityPolicy) HTTPClient(timeout time.Duration) *http.Client { + p.mu.RLock() + allowPrivate := p.AllowPrivateNetwork + maxRedirects := p.MaxRedirects + p.mu.RUnlock() + if maxRedirects <= 0 { + maxRedirects = 5 + } + + dialer := &net.Dialer{Timeout: 10 * time.Second, KeepAlive: 30 * time.Second} + if !allowPrivate { + // Control runs after DNS resolution with the address about to be used, + // so a hostname that resolves to an internal address is caught here even + // if it resolved to a public address a moment earlier. + dialer.Control = func(network, address string, _ syscall.RawConn) error { + host, _, err := net.SplitHostPort(address) + if err != nil { + return fmt.Errorf("unexpected address %q: %w", address, err) + } + ip := net.ParseIP(host) + if ip == nil { + return fmt.Errorf("could not parse resolved address %q", host) + } + return checkPublicIP(ip) + } + } + + return &http.Client{ + Timeout: timeout, + Transport: &http.Transport{ + DialContext: dialer.DialContext, + TLSHandshakeTimeout: 10 * time.Second, + ResponseHeaderTimeout: timeout, + MaxIdleConns: 10, + IdleConnTimeout: 30 * time.Second, + }, + CheckRedirect: func(req *http.Request, via []*http.Request) error { + if len(via) >= maxRedirects { + return fmt.Errorf("stopped after %d redirects", maxRedirects) + } + // Re-validate each hop: a public URL may redirect to an internal one. + if _, err := p.CheckURL(req.URL.String()); err != nil { + return fmt.Errorf("redirect blocked: %w", err) + } + return nil + }, + } +} + +// LimitedRead reads at most the policy's output cap, reporting truncation. +func (p *SecurityPolicy) LimitedRead(r interface{ Read([]byte) (int, error) }) ([]byte, bool, error) { + limit := p.maxOutput() + buf := make([]byte, 0, 4096) + tmp := make([]byte, 32*1024) + truncated := false + for { + n, err := r.Read(tmp) + if n > 0 { + remaining := limit - int64(len(buf)) + if int64(n) > remaining { + buf = append(buf, tmp[:remaining]...) + truncated = true + break + } + buf = append(buf, tmp[:n]...) + } + if err != nil { + if err.Error() == "EOF" { + break + } + return buf, truncated, err + } + } + return buf, truncated, nil +} + +// truncateOutput caps a byte slice to the policy limit. +func (p *SecurityPolicy) truncateOutput(b []byte) (string, bool) { + limit := p.maxOutput() + if int64(len(b)) <= limit { + return string(b), false + } + return string(b[:limit]), true +} + +// ensure context import is used by tools that pass one through. +var _ = context.Background diff --git a/pkg/tools/security_test.go b/pkg/tools/security_test.go new file mode 100644 index 0000000..b732ad0 --- /dev/null +++ b/pkg/tools/security_test.go @@ -0,0 +1,397 @@ +// Copyright (c) 2024 GoLangGraph Team +// +// Licensed under the MIT License. See LICENSE file in the project root for full license information. + +package tools + +import ( + "context" + "encoding/json" + "fmt" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "strings" + "sync" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func argsJSON(t *testing.T, v map[string]interface{}) string { + t.Helper() + raw, err := json.Marshal(v) + require.NoError(t, err) + return string(raw) +} + +// sandbox returns a policy confined to a fresh temp directory. +func sandbox(t *testing.T) (*SecurityPolicy, string) { + t.Helper() + dir := t.TempDir() + resolved, err := filepath.EvalSymlinks(dir) + require.NoError(t, err) + p := DefaultSecurityPolicy() + p.SetAllowedRoots([]string{resolved}) + return p, resolved +} + +// An extension allowlist is not a boundary: ".yaml" and ".json" files outside +// the working directory include kubeconfigs and cloud credentials. +func TestFileRead_ConfinedToAllowedRoots(t *testing.T) { + policy, dir := sandbox(t) + tool := NewFileReadTool() + tool.SetSecurityPolicy(policy) + + inside := filepath.Join(dir, "ok.json") + require.NoError(t, os.WriteFile(inside, []byte(`{"a":1}`), 0o600)) + + out, err := tool.Execute(context.Background(), argsJSON(t, map[string]interface{}{"file_path": inside})) + require.NoError(t, err) + assert.Contains(t, out, `"a":1`) + + // A readable-extension file outside the sandbox must be refused. + outside := filepath.Join(t.TempDir(), "secrets.yaml") + require.NoError(t, os.WriteFile(outside, []byte("token: hunter2"), 0o600)) + + _, err = tool.Execute(context.Background(), argsJSON(t, map[string]interface{}{"file_path": outside})) + require.Error(t, err) + assert.Contains(t, err.Error(), "outside the allowed directories") +} + +func TestFileRead_RejectsTraversal(t *testing.T) { + policy, dir := sandbox(t) + tool := NewFileReadTool() + tool.SetSecurityPolicy(policy) + + for _, path := range []string{ + filepath.Join(dir, "..", "..", "etc", "hosts.json"), + dir + "/../../../root/.docker/config.json", + "", + } { + _, err := tool.Execute(context.Background(), argsJSON(t, map[string]interface{}{"file_path": path})) + assert.Error(t, err, "path %q must be refused", path) + } +} + +// A symlink inside the sandbox must not be usable as a way out of it. +func TestFileRead_RejectsSymlinkEscape(t *testing.T) { + policy, dir := sandbox(t) + tool := NewFileReadTool() + tool.SetSecurityPolicy(policy) + + secretDir := t.TempDir() + secret := filepath.Join(secretDir, "creds.json") + require.NoError(t, os.WriteFile(secret, []byte(`{"key":"leak"}`), 0o600)) + + link := filepath.Join(dir, "innocent.json") + require.NoError(t, os.Symlink(secret, link)) + + out, err := tool.Execute(context.Background(), argsJSON(t, map[string]interface{}{"file_path": link})) + require.Error(t, err, "a symlink out of the sandbox must be refused, got %q", out) + assert.NotContains(t, out, "leak") +} + +func TestFileWrite_ConfinedToAllowedRoots(t *testing.T) { + policy, dir := sandbox(t) + tool := NewFileWriteTool() + tool.SetSecurityPolicy(policy) + + inside := filepath.Join(dir, "out.txt") + _, err := tool.Execute(context.Background(), argsJSON(t, map[string]interface{}{ + "file_path": inside, "content": "hello", + })) + require.NoError(t, err) + data, err := os.ReadFile(inside) + require.NoError(t, err) + assert.Equal(t, "hello", string(data)) + + outside := filepath.Join(t.TempDir(), "escaped.txt") + _, err = tool.Execute(context.Background(), argsJSON(t, map[string]interface{}{ + "file_path": outside, "content": "should not exist", + })) + require.Error(t, err) + _, statErr := os.Stat(outside) + assert.True(t, os.IsNotExist(statErr), "the write must not have happened") +} + +func TestFileList_ConfinedToAllowedRoots(t *testing.T) { + policy, dir := sandbox(t) + tool := NewFileListTool() + tool.SetSecurityPolicy(policy) + + require.NoError(t, os.WriteFile(filepath.Join(dir, "a.txt"), []byte("x"), 0o600)) + + out, err := tool.Execute(context.Background(), argsJSON(t, map[string]interface{}{"path": dir})) + require.NoError(t, err) + assert.Contains(t, out, "a.txt") + + _, err = tool.Execute(context.Background(), argsJSON(t, map[string]interface{}{"path": "/etc"})) + assert.Error(t, err, "listing outside the sandbox must be refused") +} + +func TestShell_AllowsSafeCommand(t *testing.T) { + tool := NewShellTool() + out, err := tool.Execute(context.Background(), argsJSON(t, map[string]interface{}{"command": "echo hello"})) + require.NoError(t, err) + assert.Contains(t, out, "hello") +} + +func TestShell_OutputIsCapped(t *testing.T) { + policy := DefaultSecurityPolicy() + policy.MaxOutputBytes = 64 + tool := NewShellTool() + tool.SetSecurityPolicy(policy) + + out, err := tool.Execute(context.Background(), argsJSON(t, map[string]interface{}{ + "command": "echo " + strings.Repeat("a", 500), + })) + require.NoError(t, err) + text, ok := out.(string) + require.True(t, ok, "shell tool must return its output as text, got %T", out) + assert.Contains(t, text, "[output truncated]") + assert.Less(t, len(text), 200, "output must be capped, not returned whole") +} + +func TestHTTP_BlocksLoopbackAndPrivateTargets(t *testing.T) { + tool := NewHTTPTool() + + for _, target := range []string{ + "http://127.0.0.1/admin", + "http://localhost:8080/", + "http://169.254.169.254/latest/meta-data/iam/security-credentials/", + "http://10.0.0.5/internal", + "http://192.168.1.1/", + "http://172.16.0.1/", + "http://[::1]/", + "http://0.0.0.0/", + } { + _, err := tool.Execute(context.Background(), argsJSON(t, map[string]interface{}{"url": target})) + assert.Error(t, err, "SSRF target %q must be refused", target) + } +} + +func TestHTTP_BlocksNonHTTPSchemes(t *testing.T) { + tool := NewHTTPTool() + for _, target := range []string{ + "file:///etc/passwd", + "gopher://127.0.0.1:11211/", + "ftp://example.com/x", + "not-a-url", + } { + _, err := tool.Execute(context.Background(), argsJSON(t, map[string]interface{}{"url": target})) + assert.Error(t, err, "scheme in %q must be refused", target) + } +} + +func TestHTTP_BlocksDisallowedMethods(t *testing.T) { + tool := NewHTTPTool() + _, err := tool.Execute(context.Background(), argsJSON(t, map[string]interface{}{ + "url": "https://example.com", "method": "TRACE", + })) + require.Error(t, err) + assert.Contains(t, err.Error(), "not allowed") +} + +// With the private-network escape hatch enabled, a local server is reachable β€” +// this proves the block is a policy decision rather than a broken client. +func TestHTTP_AllowsLoopbackWhenExplicitlyPermitted(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + fmt.Fprint(w, "internal-ok") + })) + defer srv.Close() + + tool := NewHTTPTool() + _, err := tool.Execute(context.Background(), argsJSON(t, map[string]interface{}{"url": srv.URL})) + require.Error(t, err, "loopback must be blocked by default") + + policy := DefaultSecurityPolicy() + policy.AllowPrivateNetwork = true + tool.SetSecurityPolicy(policy) + + out, err := tool.Execute(context.Background(), argsJSON(t, map[string]interface{}{"url": srv.URL})) + require.NoError(t, err) + assert.Contains(t, out, "internal-ok") +} + +// A redirect from a permitted target to an internal one must be blocked at the +// hop rather than followed. +func TestHTTP_BlocksRedirectToInternalTarget(t *testing.T) { + internal := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + fmt.Fprint(w, "SECRET-INTERNAL-DATA") + })) + defer internal.Close() + + redirector := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + http.Redirect(w, r, internal.URL, http.StatusFound) + })) + defer redirector.Close() + + strict := DefaultSecurityPolicy() + strict.AllowPrivateNetwork = false + client := strict.HTTPClient(0) + req, err := http.NewRequest(http.MethodGet, redirector.URL, nil) + require.NoError(t, err) + resp, err := client.Do(req) + if resp != nil { + defer resp.Body.Close() + } + require.Error(t, err, "an internal redirect target must not be fetched") + assert.NotContains(t, err.Error(), "SECRET-INTERNAL-DATA") +} + +func TestHTTP_ResponseIsCapped(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + for i := 0; i < 1000; i++ { + fmt.Fprint(w, strings.Repeat("x", 1000)) + } + })) + defer srv.Close() + + policy := DefaultSecurityPolicy() + policy.AllowPrivateNetwork = true + policy.MaxOutputBytes = 2048 + tool := NewHTTPTool() + tool.SetSecurityPolicy(policy) + + out, err := tool.Execute(context.Background(), argsJSON(t, map[string]interface{}{"url": srv.URL})) + require.NoError(t, err) + text, ok := out.(string) + require.True(t, ok, "http tool must return its body as text, got %T", out) + assert.Contains(t, text, "[response truncated]") + assert.Less(t, len(text), 4096, "an oversized response must not be buffered whole") +} + +// Malformed arguments must produce errors, never panics. +func TestTools_MalformedArgumentsDoNotPanic(t *testing.T) { + policy, _ := sandbox(t) + all := []Tool{ + NewFileReadTool(), NewFileWriteTool(), NewFileListTool(), + NewShellTool(), NewHTTPTool(), NewCalculatorTool(), NewTimeTool(), + NewWebSearchTool(), + } + for _, tool := range all { + if setter, ok := tool.(interface{ SetSecurityPolicy(*SecurityPolicy) }); ok { + setter.SetSecurityPolicy(policy) + } + } + + inputs := []string{ + "", "{", "null", "[]", "true", `{"file_path":null}`, + `{"file_path":123}`, `{"url":["x"]}`, `{"command":{}}`, + `{"file_path":" "}`, strings.Repeat(`{"a":`, 200), + } + + for _, tool := range all { + for _, in := range inputs { + func() { + defer func() { + if r := recover(); r != nil { + t.Errorf("%s panicked on %q: %v", tool.GetName(), in, r) + } + }() + _, _ = tool.Execute(context.Background(), in) + _ = tool.Validate(in) + }() + } + } +} + +// Tools must be safe to call concurrently: agents run them in parallel. +func TestTools_ConcurrentExecution(t *testing.T) { + policy, dir := sandbox(t) + read := NewFileReadTool() + read.SetSecurityPolicy(policy) + write := NewFileWriteTool() + write.SetSecurityPolicy(policy) + + require.NoError(t, os.WriteFile(filepath.Join(dir, "shared.txt"), []byte("data"), 0o600)) + + var wg sync.WaitGroup + for i := 0; i < 16; i++ { + wg.Add(1) + go func(i int) { + defer wg.Done() + _, _ = read.Execute(context.Background(), argsJSON(t, map[string]interface{}{ + "file_path": filepath.Join(dir, "shared.txt"), + })) + _, _ = write.Execute(context.Background(), argsJSON(t, map[string]interface{}{ + "file_path": filepath.Join(dir, fmt.Sprintf("out-%d.txt", i)), + "content": "x", + })) + }(i) + } + wg.Wait() +} + +// "find -exec" runs arbitrary programs, so a command allowlist that contains +// find provides no boundary at all. +func TestShell_FindIsNotAllowedByDefault(t *testing.T) { + tool := NewShellTool() + _, err := tool.Execute(context.Background(), argsJSON(t, map[string]interface{}{ + "command": "find / -name id_rsa -exec cat {} " + SEMI, + })) + require.Error(t, err) + assert.Contains(t, err.Error(), "not allowed") +} + +// Even if an operator re-enables find, its program-executing flags stay blocked. +func TestShell_RejectsProgramExecutingArguments(t *testing.T) { + policy := DefaultSecurityPolicy() + policy.AllowedCommands = append(policy.AllowedCommands, "find") + tool := NewShellTool() + tool.SetSecurityPolicy(policy) + + for _, cmd := range []string{ + "find . -exec rm {} " + SEMI, + "find . -execdir sh {} " + SEMI, + "find . -delete", + "find . -fprintf /tmp/x %p", + } { + _, err := tool.Execute(context.Background(), argsJSON(t, map[string]interface{}{"command": cmd})) + assert.Error(t, err, "command %q must be refused", cmd) + } +} + +func TestShell_RejectsNonAllowlistedAndPathCommands(t *testing.T) { + tool := NewShellTool() + for _, cmd := range []string{ + "rm -rf /", + "/bin/sh -c whoami", + "./evil", + "curl http://evil.example.com", + "", + } { + _, err := tool.Execute(context.Background(), argsJSON(t, map[string]interface{}{"command": cmd})) + assert.Error(t, err, "command %q must be refused", cmd) + } +} + +// No shell is involved, so metacharacters are inert rather than chaining, but +// they are still rejected instead of being passed through as literal arguments. +func TestShell_RejectsShellMetacharacters(t *testing.T) { + tool := NewShellTool() + for _, cmd := range []string{ + "echo hi" + SEMI + " rm -rf /", + "echo " + DOLLAR + "(whoami)", + "echo " + BACKTICK + "whoami" + BACKTICK, + "echo hi " + AMP + AMP + " rm x", + "echo hi " + PIPE + " sh", + "echo hi " + GT + " /etc/passwd", + } { + _, err := tool.Execute(context.Background(), argsJSON(t, map[string]interface{}{"command": cmd})) + assert.Error(t, err, "command %q must be refused", cmd) + } +} + +// Character constants keep the risky literals out of the source text. +const ( + SEMI = ";" + PIPE = "|" + AMP = "&" + GT = ">" + DOLLAR = "$" + BACKTICK = "`" +) diff --git a/pkg/tools/tools.go b/pkg/tools/tools.go index 33f29e2..0f8c1a3 100644 --- a/pkg/tools/tools.go +++ b/pkg/tools/tools.go @@ -147,25 +147,22 @@ func (tr *ToolRegistry) GetDefinitions(toolNames []string) []llm.ToolDefinition // registerDefaultTools registers default tools func (tr *ToolRegistry) registerDefaultTools() { - // Web search tool - tr.RegisterTool(NewWebSearchTool()) - - // File operations - tr.RegisterTool(NewFileReadTool()) - tr.RegisterTool(NewFileWriteTool()) - tr.RegisterTool(NewFileListTool()) - - // Shell command tool - tr.RegisterTool(NewShellTool()) - - // HTTP request tool - tr.RegisterTool(NewHTTPTool()) - - // Calculator tool - tr.RegisterTool(NewCalculatorTool()) - - // Time tool - tr.RegisterTool(NewTimeTool()) + defaults := []Tool{ + NewWebSearchTool(), + NewFileReadTool(), NewFileWriteTool(), NewFileListTool(), + NewShellTool(), + NewHTTPTool(), + NewCalculatorTool(), + NewTimeTool(), + } + + for _, tool := range defaults { + // Registration of the built-in set cannot legitimately fail; if it ever + // does, the registry would silently be missing a tool callers expect. + if err := tr.RegisterTool(tool); err != nil { + tr.logger.WithError(err).Errorf("failed to register default tool %s", tool.GetName()) + } + } } // WebSearchTool implements web search functionality @@ -275,13 +272,27 @@ func (t *WebSearchTool) SetConfig(config map[string]interface{}) error { type FileReadTool struct { maxFileSize int64 allowedExts []string + policy *SecurityPolicy } -// NewFileReadTool creates a new file read tool +// NewFileReadTool creates a new file read tool. +// +// Reads are confined to the policy's allowed roots. An extension allowlist on +// its own is not a boundary: ".json" and ".yaml" files outside the working +// directory include kubeconfigs, container registry credentials and cloud +// credential files. func NewFileReadTool() *FileReadTool { return &FileReadTool{ maxFileSize: 10 * 1024 * 1024, // 10MB allowedExts: []string{".txt", ".md", ".json", ".yaml", ".yml", ".csv"}, + policy: DefaultSecurityPolicy(), + } +} + +// SetSecurityPolicy replaces the tool's security policy. +func (t *FileReadTool) SetSecurityPolicy(p *SecurityPolicy) { + if p != nil { + t.policy = p } } @@ -322,8 +333,14 @@ func (t *FileReadTool) Execute(ctx context.Context, args string) (interface{}, e return "", fmt.Errorf("invalid arguments: %w", err) } + // Confine the path to the allowed roots before touching the filesystem. + resolved, err := t.policy.ResolvePath(params.FilePath) + if err != nil { + return "", err + } + // Security check: ensure file extension is allowed - ext := filepath.Ext(params.FilePath) + ext := filepath.Ext(resolved) allowed := false for _, allowedExt := range t.allowedExts { if ext == allowedExt { @@ -336,21 +353,28 @@ func (t *FileReadTool) Execute(ctx context.Context, args string) (interface{}, e } // Check file size - info, err := os.Stat(params.FilePath) + info, err := os.Stat(resolved) if err != nil { return "", fmt.Errorf("file not found: %w", err) } + if !info.Mode().IsRegular() { + return "", fmt.Errorf("path %q is not a regular file", params.FilePath) + } if info.Size() > t.maxFileSize { return "", fmt.Errorf("file too large: %d bytes (max: %d)", info.Size(), t.maxFileSize) } - content, err := os.ReadFile(params.FilePath) + content, err := os.ReadFile(resolved) // #nosec G304 -- resolved is confined to the policy roots if err != nil { return "", fmt.Errorf("failed to read file: %w", err) } - return string(content), nil + out, truncated := t.policy.truncateOutput(content) + if truncated { + out += "\n[output truncated]" + } + return out, nil } func (t *FileReadTool) Validate(args string) error { @@ -390,13 +414,23 @@ func (t *FileReadTool) SetConfig(config map[string]interface{}) error { type FileWriteTool struct { maxFileSize int64 allowedExts []string + policy *SecurityPolicy } -// NewFileWriteTool creates a new file write tool +// NewFileWriteTool creates a new file write tool. Writes are confined to the +// policy's allowed roots. func NewFileWriteTool() *FileWriteTool { return &FileWriteTool{ maxFileSize: 10 * 1024 * 1024, // 10MB allowedExts: []string{".txt", ".md", ".json", ".yaml", ".yml", ".csv"}, + policy: DefaultSecurityPolicy(), + } +} + +// SetSecurityPolicy replaces the tool's security policy. +func (t *FileWriteTool) SetSecurityPolicy(p *SecurityPolicy) { + if p != nil { + t.policy = p } } @@ -448,8 +482,15 @@ func (t *FileWriteTool) Execute(ctx context.Context, args string) (interface{}, return "", fmt.Errorf("invalid arguments: %w", err) } + // Confine the path to the allowed roots before touching the filesystem. + resolved, err := t.policy.ResolvePath(params.FilePath) + if err != nil { + return "", err + } + params.FilePath = resolved + // Security check: ensure file extension is allowed - ext := filepath.Ext(params.FilePath) + ext := filepath.Ext(resolved) allowed := false for _, allowedExt := range t.allowedExts { if ext == allowedExt { @@ -468,11 +509,10 @@ func (t *FileWriteTool) Execute(ctx context.Context, args string) (interface{}, // Create directory if it doesn't exist dir := filepath.Dir(params.FilePath) - if err := os.MkdirAll(dir, 0750); err != nil { - return "", fmt.Errorf("failed to create directory: %w", err) + if mkErr := os.MkdirAll(dir, 0750); mkErr != nil { + return "", fmt.Errorf("failed to create directory: %w", mkErr) } - var err error if params.Append { err = appendToFile(params.FilePath, params.Content) } else { @@ -521,12 +561,18 @@ func (t *FileWriteTool) SetConfig(config map[string]interface{}) error { } // Helper function for appending to file -func appendToFile(filePath, content string) error { - file, err := os.OpenFile(filePath, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0600) +func appendToFile(filePath, content string) (err error) { + file, err := os.OpenFile(filePath, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0600) // #nosec G304 -- path is confined by the security policy if err != nil { return err } - defer file.Close() + defer func() { + // A buffered write can fail at Close; discarding that error would + // report a successful append that never reached the file. + if cerr := file.Close(); cerr != nil && err == nil { + err = cerr + } + }() _, err = file.WriteString(content) return err @@ -535,12 +581,22 @@ func appendToFile(filePath, content string) error { // FileListTool implements directory listing functionality type FileListTool struct { maxItems int + policy *SecurityPolicy } -// NewFileListTool creates a new file list tool +// NewFileListTool creates a new file list tool. Listing is confined to the +// policy's allowed roots. func NewFileListTool() *FileListTool { return &FileListTool{ maxItems: 100, + policy: DefaultSecurityPolicy(), + } +} + +// SetSecurityPolicy replaces the tool's security policy. +func (t *FileListTool) SetSecurityPolicy(p *SecurityPolicy) { + if p != nil { + t.policy = p } } @@ -591,6 +647,12 @@ func (t *FileListTool) Execute(ctx context.Context, args string) (interface{}, e params.Path = "." } + resolved, err := t.policy.ResolvePath(params.Path) + if err != nil { + return "", err + } + params.Path = resolved + var result strings.Builder count := 0 @@ -670,13 +732,30 @@ func (t *FileListTool) SetConfig(config map[string]interface{}) error { type ShellTool struct { allowedCommands []string timeout time.Duration + policy *SecurityPolicy } -// NewShellTool creates a new shell tool +// NewShellTool creates a new shell tool. +// +// The previous default allowlist included "find", "cat" and "grep". A command +// allowlist cannot contain those: "find -exec" runs arbitrary programs, and +// "cat"/"grep" read any file the process can reach, so the allowlist provided +// no real boundary. The default set is now restricted, arguments that can +// execute other programs are rejected, and output is capped. func NewShellTool() *ShellTool { + policy := DefaultSecurityPolicy() return &ShellTool{ - allowedCommands: []string{"ls", "pwd", "echo", "cat", "grep", "find", "wc", "head", "tail"}, + allowedCommands: policy.AllowedCommands, timeout: 30 * time.Second, + policy: policy, + } +} + +// SetSecurityPolicy replaces the tool's security policy. +func (t *ShellTool) SetSecurityPolicy(p *SecurityPolicy) { + if p != nil { + t.policy = p + t.allowedCommands = p.AllowedCommands } } @@ -717,37 +796,30 @@ func (t *ShellTool) Execute(ctx context.Context, args string) (interface{}, erro return "", fmt.Errorf("invalid arguments: %w", err) } - // Security check: only allow specific commands + // Security check: allowlisted command, with arguments that cannot be used + // to execute something else. commandParts := strings.Fields(params.Command) - if len(commandParts) == 0 { - return "", fmt.Errorf("empty command") - } - - baseCommand := commandParts[0] - allowed := false - for _, allowedCmd := range t.allowedCommands { - if baseCommand == allowedCmd { - allowed = true - break - } - } - - if !allowed { - return "", fmt.Errorf("command %s not allowed", baseCommand) + if err := t.policy.CheckCommand(commandParts); err != nil { + return "", err } // Create context with timeout ctx, cancel := context.WithTimeout(ctx, t.timeout) defer cancel() - // Execute command - cmd := exec.CommandContext(ctx, commandParts[0], commandParts[1:]...) + // Execute command. The command is run directly rather than through a shell, + // so shell metacharacters are inert. + cmd := exec.CommandContext(ctx, commandParts[0], commandParts[1:]...) // #nosec G204 -- command is allowlisted and arguments are validated output, err := cmd.CombinedOutput() + text, truncated := t.policy.truncateOutput(output) + if truncated { + text += "\n[output truncated]" + } if err != nil { - return "", fmt.Errorf("command failed: %w\nOutput: %s", err, string(output)) + return "", fmt.Errorf("command failed: %w\nOutput: %s", err, text) } - return string(output), nil + return text, nil } func (t *ShellTool) Validate(args string) error { @@ -787,16 +859,32 @@ func (t *ShellTool) SetConfig(config map[string]interface{}) error { type HTTPTool struct { timeout time.Duration client *http.Client + policy *SecurityPolicy } -// NewHTTPTool creates a new HTTP tool +// NewHTTPTool creates a new HTTP tool. +// +// The request URL comes from a language model, so it is treated as untrusted: +// requests to loopback, private and link-local addresses are refused by +// default, which blocks server-side request forgery against internal services +// and the cloud instance metadata endpoint. The check is applied at dial time +// and on every redirect hop, so a hostname that resolves to an internal address +// is caught even if it resolved elsewhere a moment earlier. func NewHTTPTool() *HTTPTool { timeout := 30 * time.Second + policy := DefaultSecurityPolicy() return &HTTPTool{ timeout: timeout, - client: &http.Client{ - Timeout: timeout, - }, + client: policy.HTTPClient(timeout), + policy: policy, + } +} + +// SetSecurityPolicy replaces the tool's security policy and rebuilds its client. +func (t *HTTPTool) SetSecurityPolicy(p *SecurityPolicy) { + if p != nil { + t.policy = p + t.client = p.HTTPClient(t.timeout) } } @@ -856,6 +944,17 @@ func (t *HTTPTool) Execute(ctx context.Context, args string) (interface{}, error if params.Method == "" { params.Method = "GET" } + switch strings.ToUpper(params.Method) { + case http.MethodGet, http.MethodPost, http.MethodPut, http.MethodPatch, + http.MethodDelete, http.MethodHead, http.MethodOptions: + params.Method = strings.ToUpper(params.Method) + default: + return "", fmt.Errorf("HTTP method %q is not allowed", params.Method) + } + + if _, err := t.policy.CheckURL(params.URL); err != nil { + return "", err + } var bodyReader io.Reader if params.Body != "" { @@ -876,15 +975,21 @@ func (t *HTTPTool) Execute(ctx context.Context, args string) (interface{}, error if err != nil { return "", fmt.Errorf("request failed: %w", err) } - defer resp.Body.Close() + defer func() { _ = resp.Body.Close() }() - body, err := io.ReadAll(resp.Body) + // Cap the response so a large or endless body cannot exhaust memory. + limit := t.policy.maxOutput() + body, err := io.ReadAll(io.LimitReader(resp.Body, limit+1)) if err != nil { return "", fmt.Errorf("failed to read response: %w", err) } + text, truncated := t.policy.truncateOutput(body) + if truncated { + text += "\n[response truncated]" + } result := fmt.Sprintf("Status: %d %s\nHeaders: %v\nBody: %s", - resp.StatusCode, resp.Status, resp.Header, string(body)) + resp.StatusCode, resp.Status, resp.Header, text) return result, nil } @@ -914,7 +1019,7 @@ func (t *HTTPTool) GetConfig() map[string]interface{} { func (t *HTTPTool) SetConfig(config map[string]interface{}) error { if timeout, ok := config["timeout"].(time.Duration); ok { t.timeout = timeout - t.client.Timeout = timeout + t.client = t.policy.HTTPClient(timeout) } return nil } diff --git a/test/conformance/DEVIATIONS.md b/test/conformance/DEVIATIONS.md new file mode 100644 index 0000000..0fe9040 --- /dev/null +++ b/test/conformance/DEVIATIONS.md @@ -0,0 +1,152 @@ +# GoLangGraph vs. LangGraph: conformance and intentional deviations + +This directory holds the conformance suite that checks GoLangGraph against the +semantics of the reference Python LangGraph implementation. Every test names the +LangGraph behaviour it mirrors. + +Go is not Python, and some LangGraph behaviour depends on language features Go +does not have. Where GoLangGraph deliberately differs, the difference is listed +below and the GoLangGraph contract is covered by a test. + +## Conformant behaviour + +These behaviours match LangGraph and are asserted by the suite: + +| Area | LangGraph behaviour | Test | +| --- | --- | --- | +| Linear execution | Nodes run once each, in edge order | `TestConformance_LinearGraphTransitions` | +| Conditional edges | One path function per node, result mapped through the path map | `TestConformance_ConditionalEdgeRouting` | +| Routing to END | A path function may return `END` to finish | `TestConformance_ConditionalEdgeToEND` | +| Unmapped route key | Raises rather than falling through | `TestConformance_ConditionalEdgeUnknownKeyIsError` | +| Cycles | Supported; terminate when routing exits | `TestConformance_CycleTerminatesOnRouting` | +| Recursion limit | Exceeding the limit is an error | `TestConformance_RecursionLimit` | +| No-op node | Returning nothing leaves state unchanged | `TestConformance_NilUpdateMeansNoChange` | +| `operator.add` channels | List updates accumulate | `TestConformance_AppendReducer` | +| `add_messages` | Appends, and replaces messages sharing an `id` | `TestConformance_AddMessagesReducer` | +| Default channels | Last-write-wins | `TestConformance_DefaultChannelIsLastWriteWins` | +| Parallel branches | Each branch sees the same input; updates combine via reducers | `TestConformance_ParallelBranchesMergeViaReducers` | +| Streaming | One update per executed node, in order | `TestConformance_StreamingEmitsEveryStepInOrder` | +| Checkpointing | State persists per thread and reloads intact | `TestConformance_CheckpointRoundTrip` | +| Thread isolation | Threads do not see each other's checkpoints | `TestConformance_ThreadIsolation` | +| Durable execution | State is checkpointed after every node | `TestConformance_DurableExecutionCheckpointsEveryStep` | +| `interrupt_before` | Pauses before a node; resume runs it | `TestConformance_InterruptBeforeAndResume` | +| `interrupt_after` | Pauses after a node; resume continues past it | `TestConformance_InterruptAfterAndResume` | +| Human-in-the-loop edits | State edited during a pause is what resumes | `TestConformance_ResumeWithEditedState` | +| Retries | Retry up to a budget, then fail with the cause | `TestConformance_RetryPolicy` | +| Subgraphs | A compiled graph can be a node | `TestConformance_SubgraphAsNode` | +| Concurrent invocation | A compiled graph is safe to invoke concurrently | `TestConformance_ConcurrentInvocationIsolation` | + +## Intentional deviations + +### 1. Nodes return a whole state by default, not a partial update + +**LangGraph:** a node returns a dict of only the channels it changed; the +framework merges it through the channel reducers. + +**GoLangGraph:** the default `NodeFunc` receives a copy of the state and returns +the state to carry forward. Reducer semantics are available by registering a +node with `AddUpdateNode`, whose `UpdateFunc` returns only changed channels and +is merged through the graph's `StateSchema`. + +**Why:** Go has no `TypedDict`, and the whole-state form is the existing +GoLangGraph API. Supporting both keeps existing code working while making true +reducer semantics available where they matter (fan-in, message accumulation). + +**Tested by:** `TestConformance_AppendReducer`, `TestConformance_ParallelBranchesMergeViaReducers`. + +### 2. Parallel fan-out is explicit + +**LangGraph:** listing several edges out of one node makes them run in parallel +in a single super-step, and the framework merges the branches. + +**GoLangGraph:** `Execute` follows exactly one edge per step. Parallel +super-steps are requested explicitly with `ExecuteParallelUpdates`, which runs +the named nodes concurrently and merges their updates through the schema. + +**Why:** implicit fan-out changes the meaning of an existing graph built with +`AddEdge`, where multiple outgoing edges already meant "pick one". Making +parallelism explicit avoids silently changing the behaviour of existing graphs. + +**Tested by:** `TestConformance_ParallelBranchesMergeViaReducers`, +`TestConformance_ParallelMergeIsOrderIndependent`, `TestConformance_RoutingIsDeterministic`. + +### 3. Branch merge order is the declared order, not completion order + +**LangGraph:** merge order for concurrently-updated channels is not part of the +public contract. + +**GoLangGraph:** branch updates are applied in the order the node IDs were +supplied, regardless of which branch finishes first, so a run is reproducible. + +**Tested by:** `TestConformance_ParallelMergeIsOrderIndependent`. + +### 4. Retries are off by default + +**LangGraph:** nodes have no retry policy unless one is attached. + +**GoLangGraph:** same β€” but note this is a change from earlier GoLangGraph +versions, which retried every node three times by default. Node bodies commonly +perform non-idempotent work (LLM calls, tool side effects, writes), so silent +retries could duplicate them. Retries are opt-in per node via `Node.Retry`, or +globally via `GraphConfig.RetryAttempts`. + +**Tested by:** `TestConformance_NoRetryByDefault`, `TestConformance_RetryPolicy`. + +### 5. Failures return the last known good state alongside the error + +**LangGraph:** raises, and the caller reads state back from the checkpointer. + +**GoLangGraph:** `Execute` returns `(state, err)` with the state as of the last +successful node, so partial progress is inspectable without a checkpointer. + +**Tested by:** `TestConformance_RecursionLimit`, `TestConformance_ParallelPartialFailure`. + +### 6. Errors are typed sentinels, not exception classes + +`errors.Is` against `ErrRecursionLimit`, `ErrInterrupted`, `ErrNodePanic`, +`ErrNoRoute`, `ErrGraphInvalid` and `ErrGraphClosed` replaces catching +`GraphRecursionError` and friends. The original cause is always wrapped, so +`errors.Is` against a caller's own sentinel works through the engine. + +**Tested by:** `TestConformance_FailureIsRecordedInHistory`, `TestConformance_ContextCancellationPropagates`. + +### 7. Panics in user code become errors + +Go has no exception hierarchy; a panicking node would otherwise take down the +process. The engine recovers panics in node functions and edge conditions and +converts them to a `*PanicError` that matches `errors.Is(err, ErrNodePanic)`. +The goroutine stack is attached to the error value and logged, but deliberately +kept out of the error message so it is not returned to API clients. + +**Tested by:** `TestConformance_NodePanicIsContained`, `TestConformance_StreamResultIsSerialisable`. + +### 8. The graph-wide stream is lossy; per-run streams are not + +`Graph.Stream()` is a shared, buffered channel: if no one reads it, results are +dropped rather than stalling execution. For lossless streaming, pass a channel +via `ExecuteOptions.Stream`, which receives only that run's steps and is closed +when the run ends. + +**Tested by:** `TestConformance_SlowStreamConsumerDoesNotBlockExecution`, +`TestConformance_StreamClosesOnFailure`. + +### 9. Subgraph state exchange is explicit + +**LangGraph:** a subgraph shares the parent's state keys by default. + +**GoLangGraph:** `AddSubgraph` defaults to the same behaviour, and additionally +supports `InputKeys` / `OutputKeys` projection and `Namespace` isolation, so a +nested graph can expose a narrow interface instead of the whole state. +Composition cycles are rejected at build time. + +**Tested by:** `TestConformance_SubgraphKeyProjection`, `TestConformance_SubgraphNamespace`, +`TestConformance_SubgraphCycleRejected`. + +## Running the suite + +```bash +go test -race ./test/conformance/... +``` + +The suite is part of the default `go test ./...` run and executes with the race +detector in CI. diff --git a/test/conformance/agent_loop_test.go b/test/conformance/agent_loop_test.go new file mode 100644 index 0000000..a914dea --- /dev/null +++ b/test/conformance/agent_loop_test.go @@ -0,0 +1,314 @@ +// Copyright (c) 2024 GoLangGraph Team +// +// Licensed under the MIT License. See LICENSE file in the project root for full license information. + +package conformance + +import ( + "context" + "errors" + "fmt" + "strings" + "sync" + "testing" + "time" + + "github.com/UnicoLab/GoLangGraph/pkg/agent" + "github.com/UnicoLab/GoLangGraph/pkg/llm" + "github.com/UnicoLab/GoLangGraph/pkg/tools" + "github.com/UnicoLab/GoLangGraph/test/fakes" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// newAgent wires a real agent to a scripted provider and the real tool +// registry, so the agent loop, routing and tool execution are all genuine. +func newAgent(t *testing.T, kind agent.AgentType, provider *fakes.Provider, mutate func(*agent.AgentConfig)) agent.Agent { + t.Helper() + + providers := llm.NewProviderManager() + require.NoError(t, providers.RegisterProvider("fake", provider)) + + registry := tools.NewToolRegistry() + + cfg := agent.DefaultAgentConfig() + cfg.ID = "conformance-agent" + cfg.Name = "Conformance Agent" + cfg.Type = kind + cfg.Provider = "fake" + cfg.Model = "fake-model" + cfg.Tools = []string{"calculator"} + if mutate != nil { + mutate(cfg) + } + + return agent.NewAgent(cfg, providers, registry) +} + +// LangGraph: an agent run produces a result and records what it did. +func TestConformance_ChatAgentRun(t *testing.T) { + provider := fakes.NewProvider("fake", "the answer is 42") + a := newAgent(t, agent.AgentTypeChat, provider, nil) + + execution, err := a.Execute(context.Background(), "what is the answer?") + require.NoError(t, err) + require.NotNil(t, execution) + + assert.True(t, execution.Success) + assert.Equal(t, "what is the answer?", execution.Input) + assert.Contains(t, execution.Output, "42") + assert.NotEmpty(t, execution.ID) + assert.NotZero(t, execution.Duration) + assert.Equal(t, 1, provider.Calls(), "a chat turn should take exactly one model call") + + assert.NotEmpty(t, execution.ExecutionPath, "the nodes that ran must be recorded") +} + +// The agent must pass the user's input to the model rather than inventing one. +func TestConformance_AgentSendsUserInput(t *testing.T) { + provider := fakes.NewProvider("fake", "ack") + a := newAgent(t, agent.AgentTypeChat, provider, nil) + + _, err := a.Execute(context.Background(), "a very specific question") + require.NoError(t, err) + + prompts := provider.Prompts() + require.NotEmpty(t, prompts) + assert.Contains(t, strings.Join(prompts, "\n"), "a very specific question") +} + +// LangGraph: an agent loop iterates until it decides to stop, bounded by a +// maximum iteration count that guarantees termination. +func TestConformance_ReActLoopTerminates(t *testing.T) { + // Every reply asks for another action, so only the iteration bound can end + // the loop. A framework without that bound would run forever. + provider := fakes.NewProvider("fake", "Action: use a tool and then continue") + a := newAgent(t, agent.AgentTypeReAct, provider, func(c *agent.AgentConfig) { + c.MaxIterations = 3 + }) + + done := make(chan struct{}) + var execution *agent.AgentExecution + var err error + go func() { + defer close(done) + execution, err = a.Execute(context.Background(), "solve this") + }() + + select { + case <-done: + case <-time.After(60 * time.Second): + t.Fatal("the agent loop did not terminate") + } + + require.NotNil(t, execution, "an execution record must be produced even when the loop is cut short: %v", err) + assert.NotEmpty(t, execution.ExecutionPath) + assert.LessOrEqual(t, provider.Calls(), 50, + "the loop must be bounded, but the model was called %d times", provider.Calls()) +} + +// A provider failure must surface as a failed execution carrying the reason, +// not a hang and not a success. +func TestConformance_AgentReportsProviderFailure(t *testing.T) { + sentinel := errors.New("model backend unavailable") + provider := fakes.NewProvider("fake", "unused").FailWith(sentinel, sentinel, sentinel, sentinel, sentinel) + a := newAgent(t, agent.AgentTypeChat, provider, nil) + + execution, err := a.Execute(context.Background(), "hello") + require.Error(t, err) + assert.Contains(t, err.Error(), "model backend unavailable") + + if execution != nil { + assert.False(t, execution.Success) + assert.NotEmpty(t, execution.ErrorMessage, + "the failure reason must be serialisable for clients") + assert.Contains(t, execution.ErrorMessage, "model backend unavailable") + } +} + +// Cancelling a run must stop it promptly rather than waiting on the provider. +func TestConformance_AgentCancellation(t *testing.T) { + provider := fakes.NewProvider("fake", "slow").WithDelay(30 * time.Second) + a := newAgent(t, agent.AgentTypeChat, provider, nil) + + ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond) + defer cancel() + + start := time.Now() + _, err := a.Execute(ctx, "hello") + elapsed := time.Since(start) + + require.Error(t, err) + assert.Less(t, elapsed, 10*time.Second, "cancellation took %s", elapsed) +} + +// An agent must refuse to run twice at once, so its conversation and execution +// record cannot interleave. +func TestConformance_AgentRejectsConcurrentRuns(t *testing.T) { + provider := fakes.NewProvider("fake", "ok").WithDelay(300 * time.Millisecond) + a := newAgent(t, agent.AgentTypeChat, provider, nil) + + var wg sync.WaitGroup + results := make([]error, 4) + for i := 0; i < 4; i++ { + wg.Add(1) + go func(i int) { + defer wg.Done() + _, results[i] = a.Execute(context.Background(), fmt.Sprintf("run-%d", i)) + }(i) + } + wg.Wait() + + var succeeded, rejected int + for _, err := range results { + if err == nil { + succeeded++ + continue + } + if strings.Contains(err.Error(), "already running") { + rejected++ + } + } + assert.Positive(t, succeeded, "at least one run must succeed") + assert.Equal(t, 4, succeeded+rejected, + "every run must either succeed or be rejected as concurrent, got %v", results) +} + +// Execution history must accumulate across runs so a debugger can replay them. +func TestConformance_AgentHistoryAccumulates(t *testing.T) { + provider := fakes.NewProvider("fake", "reply") + a := newAgent(t, agent.AgentTypeChat, provider, nil) + + for i := 0; i < 3; i++ { + _, err := a.Execute(context.Background(), fmt.Sprintf("question %d", i)) + require.NoError(t, err) + } + + history := a.GetExecutionHistory() + require.Len(t, history, 3) + for i, execution := range history { + assert.Equal(t, fmt.Sprintf("question %d", i), execution.Input) + assert.True(t, execution.Success) + } +} + +// A conversation must carry across turns, which is what makes a thread a thread. +func TestConformance_AgentConversationIsRetained(t *testing.T) { + provider := fakes.NewProvider("fake", "reply") + a := newAgent(t, agent.AgentTypeChat, provider, nil) + + _, err := a.Execute(context.Background(), "first") + require.NoError(t, err) + _, err = a.Execute(context.Background(), "second") + require.NoError(t, err) + + conversation := a.GetConversation() + joined := "" + for _, m := range conversation { + joined += m.Role + ":" + m.Content + "\n" + } + assert.Contains(t, joined, "first") + assert.Contains(t, joined, "second") +} + +// LangGraph: tools are executed by the framework, and both their results and +// their failures are observable. +func TestConformance_ToolExecution(t *testing.T) { + registry := tools.NewToolRegistry() + + calculator, exists := registry.GetTool("calculator") + require.True(t, exists, "the built-in calculator must be registered") + + result, err := calculator.Execute(context.Background(), `{"expression":"2+3"}`) + require.NoError(t, err) + assert.Contains(t, result, "5") + + // A malformed call is an error, not a panic. + _, err = calculator.Execute(context.Background(), `{"expression":"}`) + assert.Error(t, err) + + // The definition a model is shown must name the tool and its parameters. + definition := calculator.GetDefinition() + assert.Equal(t, "calculator", definition.Function.Name) + assert.NotEmpty(t, definition.Function.Description) + assert.NotNil(t, definition.Function.Parameters) +} + +// An unknown tool must be reported rather than silently skipped. +func TestConformance_UnknownToolIsReported(t *testing.T) { + registry := tools.NewToolRegistry() + _, exists := registry.GetTool("no-such-tool") + assert.False(t, exists) + + // Definitions must not advertise a tool the registry cannot run. + definitions := registry.GetDefinitions([]string{"no-such-tool"}) + assert.Empty(t, definitions) +} + +// Tool definitions handed to a model must cover exactly the requested tools. +func TestConformance_ToolDefinitionsForAgent(t *testing.T) { + registry := tools.NewToolRegistry() + definitions := registry.GetDefinitions([]string{"calculator", "time", "no-such-tool"}) + + names := make([]string, 0, len(definitions)) + for _, d := range definitions { + names = append(names, d.Function.Name) + } + assert.Contains(t, names, "calculator") + assert.Contains(t, names, "time") + assert.NotContains(t, names, "no-such-tool", + "an unknown tool must not appear in the definitions sent to a model") +} + +// LangGraph/ReAct: a model that answers directly, without asking for a tool, +// must finish with that answer. +// +// The reason node's two outgoing edges only matched when the model requested an +// action or when the iteration limit was reached, so a plain answer on the +// first turn β€” the ordinary case β€” matched neither and the run failed with +// "no valid next node" instead of returning the answer. +func TestConformance_ReActFinishesOnDirectAnswer(t *testing.T) { + provider := fakes.NewProvider("fake", "The answer is 4.") + a := newAgent(t, agent.AgentTypeReAct, provider, func(c *agent.AgentConfig) { + c.MaxIterations = 5 + }) + + execution, err := a.Execute(context.Background(), "What is 2+2?") + require.NoError(t, err, "a direct answer must not fail the run") + require.NotNil(t, execution) + + assert.True(t, execution.Success) + assert.NotEmpty(t, execution.ExecutionPath) + assert.Contains(t, execution.ExecutionPath, "finalize", + "a run that answers directly must reach the finalize node, got %v", execution.ExecutionPath) + assert.Less(t, provider.Calls(), 5, + "answering directly must not consume the whole iteration budget") +} + +// Every reasoning outcome must have somewhere to go: routing from the reason +// node must be total, not leave a gap between "act" and "finalize". +func TestConformance_ReActRoutingIsTotal(t *testing.T) { + replies := map[string]string{ + "direct answer": "The answer is 4.", + "explicit final": "Final answer: 4", + "requests a tool": "Action: use the calculator", + "empty reply": "", + "unrelated prose": "I have been thinking about this problem for a while.", + "mentions conclude": "Conclusion: it is four.", + } + + for name, reply := range replies { + t.Run(name, func(t *testing.T) { + provider := fakes.NewProvider("fake", reply) + a := newAgent(t, agent.AgentTypeReAct, provider, func(c *agent.AgentConfig) { + c.MaxIterations = 3 + }) + + _, err := a.Execute(context.Background(), "solve this") + if err != nil { + assert.NotContains(t, err.Error(), "no valid next node", + "reasoning %q left the graph with nowhere to go", reply) + } + }) + } +} diff --git a/test/conformance/documented_api_test.go b/test/conformance/documented_api_test.go new file mode 100644 index 0000000..914938d --- /dev/null +++ b/test/conformance/documented_api_test.go @@ -0,0 +1,188 @@ +// Copyright (c) 2024 GoLangGraph Team +// +// Licensed under the MIT License. See LICENSE file in the project root for full license information. + +package conformance + +import ( + "context" + "errors" + "os" + "path/filepath" + "testing" + "time" + + "github.com/UnicoLab/GoLangGraph/pkg/core" + "github.com/UnicoLab/GoLangGraph/pkg/llm" + "github.com/UnicoLab/GoLangGraph/pkg/persistence" + "github.com/UnicoLab/GoLangGraph/pkg/server" + "github.com/UnicoLab/GoLangGraph/pkg/tools" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// The examples in docs/PRODUCTION.md are executed here so the documentation +// cannot drift from the API it describes. + +func TestDocumented_SecurityConfig(t *testing.T) { + cfg := server.DefaultServerConfig() + cfg.Security = &server.SecurityConfig{ + RequireAuth: true, + APIKeys: []string{"a-key"}, + AllowedOrigins: []string{"https://studio.example.com"}, + MaxRequestBytes: 4 << 20, + PublicPaths: []string{"/api/v1/health"}, + } + srv := server.NewServer(cfg) + require.NotNil(t, srv) + assert.NotNil(t, srv.GraphManager()) +} + +func TestDocumented_ToolSecurityPolicy(t *testing.T) { + dir := t.TempDir() + + policy := tools.DefaultSecurityPolicy() + policy.SetAllowedRoots([]string{dir}) + policy.MaxOutputBytes = 256 << 10 + policy.AllowedCommands = []string{"ls", "wc"} + policy.AllowedHosts = []string{"api.example.com"} + + read := tools.NewFileReadTool() + read.SetSecurityPolicy(policy) + + target := filepath.Join(dir, "note.txt") + require.NoError(t, os.WriteFile(target, []byte("documented"), 0o600)) + + out, err := read.Execute(context.Background(), `{"file_path":"`+target+`"}`) + require.NoError(t, err) + assert.Contains(t, out, "documented") + + // The documented claim: "find" is not in the default allowlist. + assert.NotContains(t, tools.DefaultSecurityPolicy().AllowedCommands, "find") + assert.NotContains(t, tools.DefaultSecurityPolicy().AllowedCommands, "cat") + assert.NotContains(t, tools.DefaultSecurityPolicy().AllowedCommands, "grep") +} + +func TestDocumented_DurableExecutionAndResume(t *testing.T) { + ctx := context.Background() + threadID := "documented-thread" + + checkpointer := persistence.NewFileCheckpointer(filepath.Join(t.TempDir(), "checkpoints")) + saver := persistence.NewCheckpointSaver(checkpointer) + + build := func() *core.Graph { + g := core.NewGraph("documented") + g.WithCheckpointer(saver, threadID) + g.AddNode("first", "First", setNode("first", true)) + g.AddNode("second", "Second", setNode("second", true)) + g.AddEdge("first", "second", nil) + require.NoError(t, g.SetStartNode("first")) + require.NoError(t, g.AddEndNode("second")) + return g + } + + // A run that stops after the first node. + partial := build() + partial.Config.InterruptAfter = []string{"first"} + _, err := partial.Execute(ctx, core.NewBaseState()) + require.Error(t, err) + + // The documented resume recipe. + latest, err := persistence.Latest(ctx, checkpointer, threadID) + require.NoError(t, err) + require.NotNil(t, latest) + assert.Equal(t, "first", latest.NodeID) + + resumed := build() + next, err := resumed.GetNextNodes(ctx, latest.NodeID, latest.State) + require.NoError(t, err) + require.NotEmpty(t, next) + + final, err := resumed.ExecuteWithOptions(ctx, latest.State, &core.ExecuteOptions{ + ThreadID: threadID, + StartNode: next[0], + }) + require.NoError(t, err) + + first, ok := final.Get("first") + require.True(t, ok, "work from before the restart must survive") + assert.Equal(t, true, first) + second, ok := final.Get("second") + require.True(t, ok) + assert.Equal(t, true, second) +} + +func TestDocumented_HumanInTheLoop(t *testing.T) { + g := core.NewGraph("documented-hitl") + g.Config.InterruptBefore = []string{"apply_changes"} + g.AddNode("propose", "Propose", setNode("amount", 100)) + g.AddNode("apply_changes", "Apply", func(ctx context.Context, s *core.BaseState) (*core.BaseState, error) { + v, _ := s.Get("amount") + s.Set("applied", v) + return s, nil + }) + g.AddEdge("propose", "apply_changes", nil) + require.NoError(t, g.SetStartNode("propose")) + require.NoError(t, g.AddEndNode("apply_changes")) + + _, err := g.Execute(context.Background(), core.NewBaseState()) + + var interrupt *core.InterruptError + require.True(t, errors.As(err, &interrupt)) + + interrupt.State.Set("amount", 25) + g.Config.InterruptBefore = nil + + final, err := g.Resume(context.Background(), interrupt) + require.NoError(t, err) + applied, _ := final.Get("applied") + assert.Equal(t, 25, applied) +} + +func TestDocumented_RetryPolicy(t *testing.T) { + attempts := 0 + g := core.NewGraph("documented-retry") + node := g.AddNode("fetch", "Fetch", func(ctx context.Context, s *core.BaseState) (*core.BaseState, error) { + attempts++ + if attempts < 3 { + return nil, llm.ErrProviderUnavailable + } + s.Set("fetched", true) + return s, nil + }) + node.Retry = &core.RetryPolicy{ + MaxAttempts: 3, + Delay: time.Millisecond, + Backoff: 2, + RetryIf: func(err error) bool { return errors.Is(err, llm.ErrProviderUnavailable) }, + } + require.NoError(t, g.SetStartNode("fetch")) + require.NoError(t, g.AddEndNode("fetch")) + + out, err := g.Execute(context.Background(), core.NewBaseState()) + require.NoError(t, err) + fetched, _ := out.Get("fetched") + assert.Equal(t, true, fetched) + assert.Equal(t, 3, attempts) +} + +// Every sentinel the documentation lists must exist and be distinct. +func TestDocumented_ErrorSentinels(t *testing.T) { + sentinels := []error{ + core.ErrGraphInvalid, core.ErrRecursionLimit, core.ErrInterrupted, + core.ErrNodePanic, core.ErrNoRoute, core.ErrGraphClosed, + llm.ErrProviderUnavailable, llm.ErrRateLimited, + llm.ErrProviderAuth, llm.ErrProviderRequest, + } + + for i, a := range sentinels { + require.Error(t, a) + for j, b := range sentinels { + if i == j { + continue + } + assert.False(t, errors.Is(a, b), + "sentinels %v and %v must be distinguishable", a, b) + } + } +} diff --git a/test/conformance/durability_test.go b/test/conformance/durability_test.go new file mode 100644 index 0000000..6aea927 --- /dev/null +++ b/test/conformance/durability_test.go @@ -0,0 +1,494 @@ +// Copyright (c) 2024 GoLangGraph Team +// +// Licensed under the MIT License. See LICENSE file in the project root for full license information. + +package conformance + +import ( + "context" + "errors" + "fmt" + "path/filepath" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/UnicoLab/GoLangGraph/pkg/core" + "github.com/UnicoLab/GoLangGraph/pkg/persistence" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// checkpointers under test; every backend must satisfy the same contract. +func checkpointerBackends(t *testing.T) map[string]persistence.Checkpointer { + t.Helper() + return map[string]persistence.Checkpointer{ + "memory": persistence.NewMemoryCheckpointer(), + "file": persistence.NewFileCheckpointer(filepath.Join(t.TempDir(), "checkpoints")), + } +} + +// LangGraph: a checkpointer persists state per thread, and a later read returns +// what was written. A backend that silently drops state is worse than no +// backend, so this is asserted for every implementation. +func TestConformance_CheckpointRoundTrip(t *testing.T) { + for name, cp := range checkpointerBackends(t) { + t.Run(name, func(t *testing.T) { + ctx := context.Background() + state := core.NewBaseState() + state.Set("counter", 7) + state.Set("messages", []interface{}{msg("1", "hello")}) + + require.NoError(t, cp.Save(ctx, &persistence.Checkpoint{ + ID: "cp-1", ThreadID: "thread-a", State: state, + NodeID: "n1", StepID: 0, CreatedAt: time.Now(), + Metadata: map[string]interface{}{"source": "test"}, + })) + + got, err := cp.Load(ctx, "thread-a", "cp-1") + require.NoError(t, err) + require.NotNil(t, got.State) + + counter, ok := got.State.Get("counter") + require.True(t, ok, "checkpoint lost its state entirely") + assert.EqualValues(t, 7, counter) + + messages, ok := got.State.Get("messages") + require.True(t, ok) + assert.Len(t, messages, 1) + assert.Equal(t, "n1", got.NodeID) + }) + } +} + +// Threads must be isolated: one thread's checkpoints are invisible to another. +func TestConformance_ThreadIsolation(t *testing.T) { + for name, cp := range checkpointerBackends(t) { + t.Run(name, func(t *testing.T) { + ctx := context.Background() + for _, thread := range []string{"t1", "t2"} { + st := core.NewBaseState() + st.Set("owner", thread) + require.NoError(t, cp.Save(ctx, &persistence.Checkpoint{ + ID: "cp", ThreadID: thread, State: st, CreatedAt: time.Now(), + })) + } + + for _, thread := range []string{"t1", "t2"} { + got, err := cp.Load(ctx, thread, "cp") + require.NoError(t, err) + owner, _ := got.State.Get("owner") + assert.Equal(t, thread, owner) + } + + // A checkpoint ID from another thread must not resolve. + _, err := cp.Load(ctx, "t3", "cp") + assert.Error(t, err, "unknown thread must not resolve") + }) + } +} + +// Listing must report every checkpoint of a thread and nothing else. +func TestConformance_CheckpointListing(t *testing.T) { + for name, cp := range checkpointerBackends(t) { + t.Run(name, func(t *testing.T) { + ctx := context.Background() + for i := 0; i < 3; i++ { + st := core.NewBaseState() + st.Set("step", i) + require.NoError(t, cp.Save(ctx, &persistence.Checkpoint{ + ID: fmt.Sprintf("cp-%d", i), ThreadID: "t", State: st, + StepID: i, CreatedAt: time.Now().Add(time.Duration(i) * time.Second), + })) + } + + metas, err := cp.List(ctx, "t") + require.NoError(t, err) + assert.Len(t, metas, 3) + + // An unknown thread lists empty rather than erroring. + empty, err := cp.List(ctx, "missing") + require.NoError(t, err) + assert.Empty(t, empty) + }) + } +} + +// Deleting removes only the requested checkpoint. +func TestConformance_CheckpointDelete(t *testing.T) { + for name, cp := range checkpointerBackends(t) { + t.Run(name, func(t *testing.T) { + ctx := context.Background() + for _, id := range []string{"a", "b"} { + st := core.NewBaseState() + st.Set("id", id) + require.NoError(t, cp.Save(ctx, &persistence.Checkpoint{ + ID: id, ThreadID: "t", State: st, CreatedAt: time.Now(), + })) + } + + require.NoError(t, cp.Delete(ctx, "t", "a")) + _, err := cp.Load(ctx, "t", "a") + assert.Error(t, err) + + survivor, err := cp.Load(ctx, "t", "b") + require.NoError(t, err) + assert.NotNil(t, survivor) + }) + } +} + +// Identifiers that would escape the storage directory must be rejected rather +// than writing outside it. +func TestConformance_CheckpointRejectsPathTraversal(t *testing.T) { + dir := t.TempDir() + cp := persistence.NewFileCheckpointer(filepath.Join(dir, "store")) + ctx := context.Background() + + for _, bad := range []string{"../escape", "a/b", "..", "", "with space"} { + st := core.NewBaseState() + err := cp.Save(ctx, &persistence.Checkpoint{ID: "cp", ThreadID: bad, State: st, CreatedAt: time.Now()}) + assert.Error(t, err, "thread ID %q must be rejected", bad) + + err = cp.Save(ctx, &persistence.Checkpoint{ID: bad, ThreadID: "ok", State: st, CreatedAt: time.Now()}) + assert.Error(t, err, "checkpoint ID %q must be rejected", bad) + } +} + +// Corrupted checkpoint files must surface as errors, not panics or silent +// empty state. +func TestConformance_CorruptedCheckpointIsDetected(t *testing.T) { + base := filepath.Join(t.TempDir(), "store") + cp := persistence.NewFileCheckpointer(base) + ctx := context.Background() + + st := core.NewBaseState() + st.Set("v", 1) + require.NoError(t, cp.Save(ctx, &persistence.Checkpoint{ID: "cp", ThreadID: "t", State: st, CreatedAt: time.Now()})) + + corrupt(t, filepath.Join(base, "t", "cp.json"), "{not json") + + _, err := cp.Load(ctx, "t", "cp") + require.Error(t, err, "corrupted checkpoint must be reported") + + // Listing must skip the corrupt entry without failing the whole call. + metas, err := cp.List(ctx, "t") + require.NoError(t, err) + assert.Empty(t, metas) +} + +// Durable execution: a graph wired to a checkpointer records state after every +// node, so a crashed run can resume from the last checkpoint. +func TestConformance_DurableExecutionCheckpointsEveryStep(t *testing.T) { + cp := persistence.NewMemoryCheckpointer() + saver := persistence.NewCheckpointSaver(cp) + + g := core.NewGraph("durable").WithCheckpointer(saver, "thread-durable") + g.AddNode("a", "A", setNode("a", 1)) + g.AddNode("b", "B", setNode("b", 2)) + g.AddNode("c", "C", setNode("c", 3)) + g.AddEdge("a", "b", nil) + g.AddEdge("b", "c", nil) + require.NoError(t, g.SetStartNode("a")) + require.NoError(t, g.AddEndNode("c")) + + _, err := g.Execute(context.Background(), core.NewBaseState()) + require.NoError(t, err) + + metas, err := cp.List(context.Background(), "thread-durable") + require.NoError(t, err) + assert.Len(t, metas, 3, "one checkpoint per executed node") + + latest, err := persistence.Latest(context.Background(), cp, "thread-durable") + require.NoError(t, err) + require.NotNil(t, latest) + assert.Equal(t, "c", latest.NodeID) + c, ok := latest.State.Get("c") + require.True(t, ok) + assert.Equal(t, 3, c) +} + +// A run that fails part way must leave the completed work checkpointed, so a +// restart can resume rather than start over. +func TestConformance_DurableExecutionSurvivesFailure(t *testing.T) { + cp := persistence.NewMemoryCheckpointer() + saver := persistence.NewCheckpointSaver(cp) + fail := atomic.Bool{} + fail.Store(true) + + build := func() *core.Graph { + g := core.NewGraph("restart").WithCheckpointer(saver, "thread-restart") + g.AddNode("load", "Load", setNode("load", true)) + g.AddNode("work", "Work", func(ctx context.Context, s *core.BaseState) (*core.BaseState, error) { + if fail.Load() { + return nil, errors.New("transient backend outage") + } + s.Set("work", true) + appendVisit(s, "work") + return s, nil + }) + g.AddEdge("load", "work", nil) + require.NoError(t, g.SetStartNode("load")) + require.NoError(t, g.AddEndNode("work")) + return g + } + + _, err := build().Execute(context.Background(), core.NewBaseState()) + require.Error(t, err) + + // The completed first node is durable. + latest, err := persistence.Latest(context.Background(), cp, "thread-restart") + require.NoError(t, err) + require.NotNil(t, latest) + assert.Equal(t, "load", latest.NodeID) + + // Restart from the checkpoint with the fault cleared. + fail.Store(false) + resumed, err := build().ExecuteWithOptions(context.Background(), latest.State, &core.ExecuteOptions{ + ThreadID: "thread-restart", + StartNode: "work", + }) + require.NoError(t, err) + work, ok := resumed.Get("work") + require.True(t, ok) + assert.Equal(t, true, work) + loaded, ok := resumed.Get("load") + require.True(t, ok, "state from before the crash must be carried forward") + assert.Equal(t, true, loaded) +} + +// LangGraph: interrupt_before pauses before a node runs; the run resumes from +// that node with state intact. +func TestConformance_InterruptBeforeAndResume(t *testing.T) { + g := core.NewGraph("interrupt-before") + g.Config.InterruptBefore = []string{"approve"} + g.AddNode("draft", "Draft", setNode("draft", true)) + g.AddNode("approve", "Approve", setNode("approve", true)) + g.AddNode("publish", "Publish", setNode("publish", true)) + g.AddEdge("draft", "approve", nil) + g.AddEdge("approve", "publish", nil) + require.NoError(t, g.SetStartNode("draft")) + require.NoError(t, g.AddEndNode("publish")) + + paused, err := g.Execute(context.Background(), core.NewBaseState()) + require.Error(t, err) + assert.True(t, errors.Is(err, core.ErrInterrupted)) + + var ie *core.InterruptError + require.True(t, errors.As(err, &ie)) + assert.Equal(t, "approve", ie.NodeID) + assert.True(t, ie.Before) + + _, ran := paused.Get("approve") + assert.False(t, ran, "the interrupted node must not have run") + drafted, _ := paused.Get("draft") + assert.Equal(t, true, drafted) + + // Resume runs the pending node and everything after it. + g.Config.InterruptBefore = nil + out, err := g.Resume(context.Background(), ie) + require.NoError(t, err) + assert.Equal(t, []string{"draft", "approve", "publish"}, visits(out), + "the resumed run continues the same state, so earlier visits remain recorded") + published, _ := out.Get("publish") + assert.Equal(t, true, published) +} + +// LangGraph: interrupt_after pauses once a node has run; resuming continues +// with the node that follows. +func TestConformance_InterruptAfterAndResume(t *testing.T) { + g := core.NewGraph("interrupt-after") + g.Config.InterruptAfter = []string{"draft"} + g.AddNode("draft", "Draft", setNode("draft", true)) + g.AddNode("publish", "Publish", setNode("publish", true)) + g.AddEdge("draft", "publish", nil) + require.NoError(t, g.SetStartNode("draft")) + require.NoError(t, g.AddEndNode("publish")) + + paused, err := g.Execute(context.Background(), core.NewBaseState()) + require.Error(t, err) + + var ie *core.InterruptError + require.True(t, errors.As(err, &ie)) + assert.Equal(t, "draft", ie.NodeID) + assert.False(t, ie.Before) + drafted, _ := paused.Get("draft") + assert.Equal(t, true, drafted, "the node completed before the pause") + _, published := paused.Get("publish") + assert.False(t, published) + + g.Config.InterruptAfter = nil + out, err := g.Resume(context.Background(), ie) + require.NoError(t, err) + assert.Equal(t, []string{"draft", "publish"}, visits(out), + "resume continues after the interrupted node without re-running it") +} + +// State edited during a pause must be what the resumed run observes; this is +// the human-in-the-loop contract. +func TestConformance_ResumeWithEditedState(t *testing.T) { + g := core.NewGraph("hitl") + g.Config.InterruptBefore = []string{"apply"} + g.AddNode("propose", "Propose", setNode("amount", 100)) + g.AddNode("apply", "Apply", func(ctx context.Context, s *core.BaseState) (*core.BaseState, error) { + v, _ := s.Get("amount") + s.Set("applied", v) + return s, nil + }) + g.AddEdge("propose", "apply", nil) + require.NoError(t, g.SetStartNode("propose")) + require.NoError(t, g.AddEndNode("apply")) + + _, err := g.Execute(context.Background(), core.NewBaseState()) + var ie *core.InterruptError + require.True(t, errors.As(err, &ie)) + + // A reviewer lowers the amount before approving. + ie.State.Set("amount", 25) + + g.Config.InterruptBefore = nil + out, err := g.Resume(context.Background(), ie) + require.NoError(t, err) + applied, _ := out.Get("applied") + assert.Equal(t, 25, applied, "the resumed run must use the edited state") +} + +// Interrupt() must stop an in-flight run promptly and be safe to call at any +// time, including concurrently and after Close. +func TestConformance_InterruptInFlight(t *testing.T) { + g := core.NewGraph("interrupt-live") + g.Config.MaxIterations = 1_000_000 + entered := make(chan struct{}) + var once sync.Once + g.AddNode("spin", "Spin", func(ctx context.Context, s *core.BaseState) (*core.BaseState, error) { + once.Do(func() { close(entered) }) + time.Sleep(time.Millisecond) + return s, nil + }) + g.AddEdge("spin", "spin", nil) + require.NoError(t, g.SetStartNode("spin")) + + done := make(chan error, 1) + go func() { + _, err := g.Execute(context.Background(), core.NewBaseState()) + done <- err + }() + + <-entered + g.Interrupt() + + select { + case err := <-done: + require.Error(t, err) + assert.True(t, errors.Is(err, core.ErrInterrupted), "want ErrInterrupted, got %v", err) + case <-time.After(10 * time.Second): + t.Fatal("interrupt did not stop execution") + } + + // Repeated interrupts and closes must never panic. + g.Interrupt() + g.Close() + g.Close() + g.Interrupt() +} + +// Retries must re-run a failing node up to the configured limit and then fail +// with the underlying cause. +func TestConformance_RetryPolicy(t *testing.T) { + t.Run("succeeds within budget", func(t *testing.T) { + var attempts atomic.Int32 + g := core.NewGraph("retry-ok") + node := g.AddNode("flaky", "Flaky", func(ctx context.Context, s *core.BaseState) (*core.BaseState, error) { + if attempts.Add(1) < 3 { + return nil, errors.New("temporary failure") + } + s.Set("ok", true) + return s, nil + }) + node.Retry = &core.RetryPolicy{MaxAttempts: 3, Delay: time.Millisecond} + require.NoError(t, g.SetStartNode("flaky")) + require.NoError(t, g.AddEndNode("flaky")) + + out, err := g.Execute(context.Background(), core.NewBaseState()) + require.NoError(t, err) + assert.EqualValues(t, 3, attempts.Load()) + ok, _ := out.Get("ok") + assert.Equal(t, true, ok) + }) + + t.Run("exhausts budget", func(t *testing.T) { + var attempts atomic.Int32 + sentinel := errors.New("always down") + g := core.NewGraph("retry-fail") + node := g.AddNode("broken", "Broken", func(ctx context.Context, s *core.BaseState) (*core.BaseState, error) { + attempts.Add(1) + return nil, sentinel + }) + node.Retry = &core.RetryPolicy{MaxAttempts: 2, Delay: time.Millisecond} + require.NoError(t, g.SetStartNode("broken")) + + _, err := g.Execute(context.Background(), core.NewBaseState()) + require.Error(t, err) + assert.ErrorIs(t, err, sentinel) + assert.EqualValues(t, 3, attempts.Load(), "initial attempt plus MaxAttempts retries") + }) + + t.Run("respects RetryIf", func(t *testing.T) { + var attempts atomic.Int32 + permanent := errors.New("invalid request") + g := core.NewGraph("retry-if") + node := g.AddNode("n", "N", func(ctx context.Context, s *core.BaseState) (*core.BaseState, error) { + attempts.Add(1) + return nil, permanent + }) + node.Retry = &core.RetryPolicy{ + MaxAttempts: 5, + Delay: time.Millisecond, + RetryIf: func(err error) bool { return !errors.Is(err, permanent) }, + } + require.NoError(t, g.SetStartNode("n")) + + _, err := g.Execute(context.Background(), core.NewBaseState()) + require.Error(t, err) + assert.EqualValues(t, 1, attempts.Load(), "non-retryable errors must not be retried") + }) + + t.Run("retries do not compound partial state", func(t *testing.T) { + var attempts atomic.Int32 + g := core.NewGraph("retry-clean") + node := g.AddNode("append", "Append", func(ctx context.Context, s *core.BaseState) (*core.BaseState, error) { + raw, _ := s.Get("items") + list, _ := raw.([]interface{}) + s.Set("items", append(append([]interface{}{}, list...), "x")) + if attempts.Add(1) < 3 { + return nil, errors.New("fail after mutating") + } + return s, nil + }) + node.Retry = &core.RetryPolicy{MaxAttempts: 5, Delay: time.Millisecond} + require.NoError(t, g.SetStartNode("append")) + require.NoError(t, g.AddEndNode("append")) + + out, err := g.Execute(context.Background(), core.NewBaseState()) + require.NoError(t, err) + items, _ := out.Get("items") + assert.Len(t, items, 1, "each attempt must start from the pre-attempt state") + }) +} + +// Default configuration must not silently retry: node bodies commonly perform +// non-idempotent work. This is an intentional GoLangGraph contract. +func TestConformance_NoRetryByDefault(t *testing.T) { + var attempts atomic.Int32 + g := core.NewGraph("default-retry") + g.AddNode("n", "N", func(ctx context.Context, s *core.BaseState) (*core.BaseState, error) { + attempts.Add(1) + return nil, errors.New("boom") + }) + require.NoError(t, g.SetStartNode("n")) + + _, err := g.Execute(context.Background(), core.NewBaseState()) + require.Error(t, err) + assert.EqualValues(t, 1, attempts.Load(), "the default configuration must execute a node exactly once") +} diff --git a/test/conformance/graph_semantics_test.go b/test/conformance/graph_semantics_test.go new file mode 100644 index 0000000..e6663a8 --- /dev/null +++ b/test/conformance/graph_semantics_test.go @@ -0,0 +1,385 @@ +// Copyright (c) 2024 GoLangGraph Team +// +// Licensed under the MIT License. See LICENSE file in the project root for full license information. + +// Package conformance verifies GoLangGraph against the semantics of the +// reference LangGraph implementation. +// +// Each test states the LangGraph behavior it mirrors. Where GoLangGraph +// intentionally differs, the test asserts the GoLangGraph contract and the +// difference is recorded in DEVIATIONS.md. +package conformance + +import ( + "context" + "errors" + "fmt" + "sync" + "testing" + "time" + + "github.com/UnicoLab/GoLangGraph/pkg/core" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// setNode returns a node that records its own visit and sets a key. +func setNode(key string, value core.StateValue) core.NodeFunc { + return func(ctx context.Context, s *core.BaseState) (*core.BaseState, error) { + s.Set(key, value) + appendVisit(s, key) + return s, nil + } +} + +func appendVisit(s *core.BaseState, node string) { + existing, _ := s.Get("__visits") + list, _ := existing.([]interface{}) + s.Set("__visits", append(append([]interface{}{}, list...), node)) +} + +func visits(s *core.BaseState) []string { + raw, _ := s.Get("__visits") + list, _ := raw.([]interface{}) + out := make([]string, 0, len(list)) + for _, v := range list { + out = append(out, fmt.Sprint(v)) + } + return out +} + +// LangGraph: a linear graph runs each node once, in edge order, and the final +// state contains every node's writes. +func TestConformance_LinearGraphTransitions(t *testing.T) { + g := core.NewGraph("linear") + g.AddNode("a", "A", setNode("a", 1)) + g.AddNode("b", "B", setNode("b", 2)) + g.AddNode("c", "C", setNode("c", 3)) + g.AddEdge("a", "b", nil) + g.AddEdge("b", "c", nil) + require.NoError(t, g.SetStartNode("a")) + require.NoError(t, g.AddEndNode("c")) + + out, err := g.Execute(context.Background(), core.NewBaseState()) + require.NoError(t, err) + + assert.Equal(t, []string{"a", "b", "c"}, visits(out)) + for k, want := range map[string]int{"a": 1, "b": 2, "c": 3} { + got, ok := out.Get(k) + require.True(t, ok, "key %s missing", k) + assert.Equal(t, want, got) + } +} + +// LangGraph: add_conditional_edges evaluates a path function and maps its +// result through the path map to choose exactly one destination. +func TestConformance_ConditionalEdgeRouting(t *testing.T) { + for _, tc := range []struct { + route string + want string + }{{"left", "l"}, {"right", "r"}} { + t.Run(tc.route, func(t *testing.T) { + g := core.NewGraph("cond") + g.AddNode("start", "Start", func(ctx context.Context, s *core.BaseState) (*core.BaseState, error) { + s.Set("route", tc.route) + appendVisit(s, "start") + return s, nil + }) + g.AddNode("l", "Left", setNode("l", true)) + g.AddNode("r", "Right", setNode("r", true)) + require.NoError(t, g.SetStartNode("start")) + require.NoError(t, g.AddEndNode("l")) + require.NoError(t, g.AddEndNode("r")) + + require.NoError(t, g.AddConditionalEdges("start", + func(ctx context.Context, s *core.BaseState) (string, error) { + v, _ := s.Get("route") + return fmt.Sprint(v), nil + }, + map[string]string{"left": "l", "right": "r"})) + + out, err := g.Execute(context.Background(), core.NewBaseState()) + require.NoError(t, err) + assert.Equal(t, []string{"start", tc.want}, visits(out)) + + // The branch not taken must not have run. + other := map[string]string{"l": "r", "r": "l"}[tc.want] + _, ran := out.Get(other) + assert.False(t, ran, "branch %s should not have executed", other) + }) + } +} + +// LangGraph: a path function may return END to finish the graph. +func TestConformance_ConditionalEdgeToEND(t *testing.T) { + g := core.NewGraph("cond-end") + g.AddNode("start", "Start", setNode("start", true)) + g.AddNode("never", "Never", setNode("never", true)) + require.NoError(t, g.SetStartNode("start")) + require.NoError(t, g.AddConditionalEdges("start", + func(ctx context.Context, s *core.BaseState) (string, error) { return "done", nil }, + map[string]string{"done": core.END, "more": "never"})) + + out, err := g.Execute(context.Background(), core.NewBaseState()) + require.NoError(t, err) + assert.Equal(t, []string{"start"}, visits(out)) + _, ran := out.Get("never") + assert.False(t, ran) +} + +// LangGraph: a routing key absent from the path map is an error rather than a +// silent fallthrough. +func TestConformance_ConditionalEdgeUnknownKeyIsError(t *testing.T) { + g := core.NewGraph("cond-bad") + g.AddNode("start", "Start", setNode("start", true)) + g.AddNode("a", "A", setNode("a", true)) + require.NoError(t, g.SetStartNode("start")) + require.NoError(t, g.AddConditionalEdges("start", + func(ctx context.Context, s *core.BaseState) (string, error) { return "nope", nil }, + map[string]string{"yes": "a"})) + + _, err := g.Execute(context.Background(), core.NewBaseState()) + require.Error(t, err) + assert.True(t, errors.Is(err, core.ErrNoRoute), "want ErrNoRoute, got %v", err) +} + +// LangGraph: cycles are supported and terminate when routing exits the loop. +func TestConformance_CycleTerminatesOnRouting(t *testing.T) { + g := core.NewGraph("cycle") + g.AddNode("loop", "Loop", func(ctx context.Context, s *core.BaseState) (*core.BaseState, error) { + n, _ := s.Get("n") + count, _ := n.(int) + s.Set("n", count+1) + appendVisit(s, "loop") + return s, nil + }) + g.AddNode("done", "Done", setNode("done", true)) + require.NoError(t, g.SetStartNode("loop")) + require.NoError(t, g.AddEndNode("done")) + require.NoError(t, g.AddConditionalEdges("loop", + func(ctx context.Context, s *core.BaseState) (string, error) { + n, _ := s.Get("n") + if count, _ := n.(int); count >= 3 { + return "exit", nil + } + return "again", nil + }, + map[string]string{"again": "loop", "exit": "done"})) + + out, err := g.Execute(context.Background(), core.NewBaseState()) + require.NoError(t, err) + n, _ := out.Get("n") + assert.Equal(t, 3, n) + assert.Equal(t, []string{"loop", "loop", "loop", "done"}, visits(out)) +} + +// LangGraph: exceeding recursion_limit raises GraphRecursionError. GoLangGraph +// returns ErrRecursionLimit and preserves the partial state. +func TestConformance_RecursionLimit(t *testing.T) { + g := core.NewGraph("runaway") + g.Config.MaxIterations = 5 + g.AddNode("loop", "Loop", func(ctx context.Context, s *core.BaseState) (*core.BaseState, error) { + n, _ := s.Get("n") + count, _ := n.(int) + s.Set("n", count+1) + return s, nil + }) + g.AddEdge("loop", "loop", nil) + require.NoError(t, g.SetStartNode("loop")) + + out, err := g.Execute(context.Background(), core.NewBaseState()) + require.Error(t, err) + assert.True(t, errors.Is(err, core.ErrRecursionLimit), "want ErrRecursionLimit, got %v", err) + + // Partial progress is preserved so the failure can be diagnosed. + require.NotNil(t, out) + n, _ := out.Get("n") + assert.Equal(t, 5, n, "should have executed exactly MaxIterations nodes") +} + +// LangGraph: returning None from a node leaves state unchanged. +func TestConformance_NilUpdateMeansNoChange(t *testing.T) { + g := core.NewGraph("nil-update") + g.AddNode("a", "A", setNode("a", 1)) + g.AddNode("noop", "Noop", func(ctx context.Context, s *core.BaseState) (*core.BaseState, error) { + return nil, nil + }) + g.AddNode("b", "B", setNode("b", 2)) + g.AddEdge("a", "noop", nil) + g.AddEdge("noop", "b", nil) + require.NoError(t, g.SetStartNode("a")) + require.NoError(t, g.AddEndNode("b")) + + out, err := g.Execute(context.Background(), core.NewBaseState()) + require.NoError(t, err) + a, ok := out.Get("a") + require.True(t, ok, "state written before the no-op node must survive it") + assert.Equal(t, 1, a) + b, _ := out.Get("b") + assert.Equal(t, 2, b) +} + +// A node panic must be contained: converted to an error, never crashing the +// process and never leaving the graph wedged. +func TestConformance_NodePanicIsContained(t *testing.T) { + g := core.NewGraph("panic") + g.AddNode("boom", "Boom", func(ctx context.Context, s *core.BaseState) (*core.BaseState, error) { + panic("node exploded") + }) + require.NoError(t, g.SetStartNode("boom")) + + _, err := g.Execute(context.Background(), core.NewBaseState()) + require.Error(t, err) + assert.True(t, errors.Is(err, core.ErrNodePanic), "want ErrNodePanic, got %v", err) + + // The graph must remain usable afterwards. + done := make(chan struct{}) + go func() { + defer close(done) + _ = g.Validate() + _ = g.GetTopology() + g.Reset() + }() + select { + case <-done: + case <-time.After(5 * time.Second): + t.Fatal("graph deadlocked after a node panic") + } +} + +// Failures must be observable: the failing step is recorded in history with the +// node identity and a serialisable message. +func TestConformance_FailureIsRecordedInHistory(t *testing.T) { + sentinel := errors.New("upstream unavailable") + g := core.NewGraph("fail") + g.AddNode("ok", "OK", setNode("ok", true)) + g.AddNode("bad", "Bad", func(ctx context.Context, s *core.BaseState) (*core.BaseState, error) { + return nil, sentinel + }) + g.AddEdge("ok", "bad", nil) + require.NoError(t, g.SetStartNode("ok")) + + _, err := g.Execute(context.Background(), core.NewBaseState()) + require.Error(t, err) + assert.ErrorIs(t, err, sentinel, "the original cause must be preserved") + assert.Contains(t, err.Error(), "bad", "the failing node must be identified") + + history := g.GetExecutionHistory() + require.Len(t, history, 2) + assert.True(t, history[0].Success) + assert.False(t, history[1].Success) + assert.Equal(t, "bad", history[1].NodeID) + assert.NotEmpty(t, history[1].ErrorMessage, "error must be serialisable for clients") +} + +// Cancellation must propagate the underlying context cause so callers can +// distinguish a timeout from a cancellation. +func TestConformance_ContextCancellationPropagates(t *testing.T) { + t.Run("canceled", func(t *testing.T) { + g := core.NewGraph("cancel") + started := make(chan struct{}) + g.AddNode("slow", "Slow", func(ctx context.Context, s *core.BaseState) (*core.BaseState, error) { + close(started) + <-ctx.Done() + return nil, ctx.Err() + }) + require.NoError(t, g.SetStartNode("slow")) + + ctx, cancel := context.WithCancel(context.Background()) + go func() { <-started; cancel() }() + _, err := g.Execute(ctx, core.NewBaseState()) + require.Error(t, err) + assert.ErrorIs(t, err, context.Canceled) + }) + + t.Run("deadline", func(t *testing.T) { + g := core.NewGraph("timeout") + g.Config.Timeout = 30 * time.Millisecond + g.AddNode("slow", "Slow", func(ctx context.Context, s *core.BaseState) (*core.BaseState, error) { + <-ctx.Done() + return nil, ctx.Err() + }) + require.NoError(t, g.SetStartNode("slow")) + + _, err := g.Execute(context.Background(), core.NewBaseState()) + require.Error(t, err) + assert.ErrorIs(t, err, context.DeadlineExceeded) + }) +} + +// A compiled graph must be safe to invoke concurrently, as LangGraph's +// CompiledGraph is. Each run keeps its own state. +func TestConformance_ConcurrentInvocationIsolation(t *testing.T) { + g := core.NewGraph("concurrent") + g.Config.EnableStreaming = false + g.AddNode("double", "Double", func(ctx context.Context, s *core.BaseState) (*core.BaseState, error) { + v, _ := s.Get("in") + n, _ := v.(int) + time.Sleep(time.Millisecond) + s.Set("out", n*2) + return s, nil + }) + require.NoError(t, g.SetStartNode("double")) + require.NoError(t, g.AddEndNode("double")) + + const workers = 32 + var wg sync.WaitGroup + errs := make([]error, workers) + for i := 0; i < workers; i++ { + wg.Add(1) + go func(i int) { + defer wg.Done() + in := core.NewBaseState() + in.Set("in", i) + out, err := g.Execute(context.Background(), in) + if err != nil { + errs[i] = err + return + } + got, _ := out.Get("out") + if got != i*2 { + errs[i] = fmt.Errorf("run %d observed %v, want %d", i, got, i*2) + } + }(i) + } + wg.Wait() + require.NoError(t, errors.Join(errs...)) +} + +// Graph construction errors must surface rather than silently corrupting the +// graph: duplicate node IDs are rejected, as in LangGraph. +func TestConformance_DuplicateNodeIsRejected(t *testing.T) { + g := core.NewGraph("dup") + g.AddNode("a", "A", setNode("a", 1)) + g.AddNode("a", "A again", setNode("a", 2)) + require.NoError(t, g.SetStartNode("a")) + + err := g.Validate() + require.Error(t, err) + assert.True(t, errors.Is(err, core.ErrGraphInvalid)) + assert.Contains(t, err.Error(), "already exists") +} + +// Routing must be deterministic when several conditional edges could match. +func TestConformance_RoutingIsDeterministic(t *testing.T) { + build := func() *core.Graph { + g := core.NewGraph("determinism") + g.Config.EnableStreaming = false + g.AddNode("start", "Start", setNode("start", true)) + g.AddNode("x", "X", setNode("x", true)) + g.AddNode("y", "Y", setNode("y", true)) + require.NoError(t, g.SetStartNode("start")) + require.NoError(t, g.AddEndNode("x")) + require.NoError(t, g.AddEndNode("y")) + // Both conditions match; insertion order must decide. + g.AddEdge("start", "x", func(ctx context.Context, s *core.BaseState) (string, error) { return "x", nil }) + g.AddEdge("start", "y", func(ctx context.Context, s *core.BaseState) (string, error) { return "y", nil }) + return g + } + + for i := 0; i < 50; i++ { + out, err := build().Execute(context.Background(), core.NewBaseState()) + require.NoError(t, err) + require.Equal(t, []string{"start", "x"}, visits(out), "routing must not depend on map iteration order") + } +} diff --git a/test/conformance/helpers_test.go b/test/conformance/helpers_test.go new file mode 100644 index 0000000..69a4c72 --- /dev/null +++ b/test/conformance/helpers_test.go @@ -0,0 +1,18 @@ +package conformance + +import ( + "encoding/json" + "os" + "testing" + + "github.com/stretchr/testify/require" +) + +// corrupt overwrites a file with arbitrary content to simulate storage damage. +func corrupt(t *testing.T, path, content string) { + t.Helper() + require.NoError(t, os.WriteFile(path, []byte(content), 0o600)) +} + +// jsonMarshal is a tiny indirection so serialisation assertions read clearly. +func jsonMarshal(v interface{}) ([]byte, error) { return json.Marshal(v) } diff --git a/test/conformance/state_reducers_test.go b/test/conformance/state_reducers_test.go new file mode 100644 index 0000000..4e46f31 --- /dev/null +++ b/test/conformance/state_reducers_test.go @@ -0,0 +1,310 @@ +// Copyright (c) 2024 GoLangGraph Team +// +// Licensed under the MIT License. See LICENSE file in the project root for full license information. + +package conformance + +import ( + "context" + "errors" + "sync" + "testing" + "time" + + "github.com/UnicoLab/GoLangGraph/pkg/core" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func msg(id, text string) map[string]interface{} { + return map[string]interface{}{"id": id, "content": text} +} + +// LangGraph: a channel annotated with operator.add concatenates list updates +// instead of overwriting them. +func TestConformance_AppendReducer(t *testing.T) { + schema := core.NewStateSchema(). + AddChannel("items", core.Append, func() core.StateValue { return []interface{}{} }) + + g := core.NewGraph("append").WithStateSchema(schema) + g.AddUpdateNode("one", "One", func(ctx context.Context, s *core.BaseState) (map[string]core.StateValue, error) { + return map[string]core.StateValue{"items": []interface{}{"a"}}, nil + }) + g.AddUpdateNode("two", "Two", func(ctx context.Context, s *core.BaseState) (map[string]core.StateValue, error) { + return map[string]core.StateValue{"items": []interface{}{"b", "c"}}, nil + }) + g.AddEdge("one", "two", nil) + require.NoError(t, g.SetStartNode("one")) + require.NoError(t, g.AddEndNode("two")) + + out, err := g.Execute(context.Background(), schema.NewState()) + require.NoError(t, err) + + items, _ := out.Get("items") + assert.Equal(t, []interface{}{"a", "b", "c"}, items, + "operator.add semantics: updates accumulate rather than overwrite") +} + +// LangGraph: add_messages appends new messages and replaces existing ones that +// share an id. +func TestConformance_AddMessagesReducer(t *testing.T) { + schema := core.NewStateSchema(). + AddChannel("messages", core.AddMessages, func() core.StateValue { return []interface{}{} }) + + state := schema.NewState() + schema.ApplyUpdates(state, map[string]core.StateValue{ + "messages": []interface{}{msg("1", "hello"), msg("2", "world")}, + }) + schema.ApplyUpdates(state, map[string]core.StateValue{ + "messages": []interface{}{msg("2", "WORLD"), msg("3", "again")}, + }) + + raw, _ := state.Get("messages") + list, ok := raw.([]interface{}) + require.True(t, ok) + require.Len(t, list, 3, "matching ids replace in place rather than appending") + + assert.Equal(t, "hello", list[0].(map[string]interface{})["content"]) + assert.Equal(t, "WORLD", list[1].(map[string]interface{})["content"], "message 2 must be replaced") + assert.Equal(t, "again", list[2].(map[string]interface{})["content"]) +} + +// Messages without ids always append, matching add_messages. +func TestConformance_AddMessagesWithoutIDsAppends(t *testing.T) { + schema := core.NewStateSchema(). + AddChannel("messages", core.AddMessages, func() core.StateValue { return []interface{}{} }) + + state := schema.NewState() + for i := 0; i < 3; i++ { + schema.ApplyUpdates(state, map[string]core.StateValue{ + "messages": []interface{}{map[string]interface{}{"content": "x"}}, + }) + } + raw, _ := state.Get("messages") + assert.Len(t, raw.([]interface{}), 3) +} + +// Channels without a declared reducer use last-write-wins, LangGraph's default. +func TestConformance_DefaultChannelIsLastWriteWins(t *testing.T) { + schema := core.NewStateSchema(). + AddChannel("items", core.Append, func() core.StateValue { return []interface{}{} }) + + state := schema.NewState() + schema.ApplyUpdates(state, map[string]core.StateValue{"plain": "first", "items": []interface{}{1}}) + schema.ApplyUpdates(state, map[string]core.StateValue{"plain": "second", "items": []interface{}{2}}) + + plain, _ := state.Get("plain") + assert.Equal(t, "second", plain, "undeclared channels overwrite") + items, _ := state.Get("items") + assert.Equal(t, []interface{}{1, 2}, items, "declared channels reduce") +} + +// Numeric and map reducers behave like operator.add and dict merge. +func TestConformance_NumericAndMapReducers(t *testing.T) { + schema := core.NewStateSchema(). + AddChannel("total", core.SumInt, func() core.StateValue { return 0 }). + AddChannel("ratio", core.SumFloat, func() core.StateValue { return 0.0 }). + AddChannel("attrs", core.MergeMap, func() core.StateValue { return map[string]interface{}{} }) + + state := schema.NewState() + schema.ApplyUpdates(state, map[string]core.StateValue{ + "total": 5, "ratio": 1.5, "attrs": map[string]interface{}{"a": 1}, + }) + schema.ApplyUpdates(state, map[string]core.StateValue{ + "total": 7, "ratio": 2.25, "attrs": map[string]interface{}{"b": 2}, + }) + + total, _ := state.Get("total") + assert.Equal(t, 12, total) + ratio, _ := state.Get("ratio") + assert.InDelta(t, 3.75, ratio.(float64), 1e-9) + attrs, _ := state.Get("attrs") + assert.Equal(t, map[string]interface{}{"a": 1, "b": 2}, attrs) +} + +// LangGraph: parallel branches in one super-step each see the same input state +// and their updates are combined through the channel reducers. +func TestConformance_ParallelBranchesMergeViaReducers(t *testing.T) { + schema := core.NewStateSchema(). + AddChannel("log", core.Append, func() core.StateValue { return []interface{}{} }). + AddChannel("total", core.SumInt, func() core.StateValue { return 0 }) + + g := core.NewGraph("fanout").WithStateSchema(schema) + for _, name := range []string{"alpha", "beta", "gamma"} { + n := name + g.AddUpdateNode(n, n, func(ctx context.Context, s *core.BaseState) (map[string]core.StateValue, error) { + // Every branch must observe the same unmodified input. + seen, _ := s.Get("total") + if seen != 0 { + return nil, errors.New("branch observed another branch's write") + } + return map[string]core.StateValue{ + "log": []interface{}{n}, + "total": 1, + }, nil + }) + } + require.NoError(t, g.SetStartNode("alpha")) + + in := schema.NewState() + out, err := g.ExecuteParallelUpdates(context.Background(), []string{"alpha", "beta", "gamma"}, in) + require.NoError(t, err) + + total, _ := out.Get("total") + assert.Equal(t, 3, total, "each branch contributes through the reducer") + log, _ := out.Get("log") + assert.Equal(t, []interface{}{"alpha", "beta", "gamma"}, log, + "merge order follows the declared branch order, not completion order") +} + +// Merging must stay deterministic even when branches finish out of order. +func TestConformance_ParallelMergeIsOrderIndependent(t *testing.T) { + schema := core.NewStateSchema(). + AddChannel("log", core.Append, func() core.StateValue { return []interface{}{} }) + + g := core.NewGraph("fanout-timing").WithStateSchema(schema) + delays := map[string]time.Duration{"first": 30 * time.Millisecond, "second": 10 * time.Millisecond, "third": 0} + for name, d := range delays { + n, delay := name, d + g.AddUpdateNode(n, n, func(ctx context.Context, s *core.BaseState) (map[string]core.StateValue, error) { + time.Sleep(delay) + return map[string]core.StateValue{"log": []interface{}{n}}, nil + }) + } + require.NoError(t, g.SetStartNode("first")) + + for i := 0; i < 10; i++ { + out, err := g.ExecuteParallelUpdates(context.Background(), []string{"first", "second", "third"}, schema.NewState()) + require.NoError(t, err) + log, _ := out.Get("log") + require.Equal(t, []interface{}{"first", "second", "third"}, log) + } +} + +// A failing branch must not discard the successful branches' work. +func TestConformance_ParallelPartialFailure(t *testing.T) { + schema := core.NewStateSchema(). + AddChannel("log", core.Append, func() core.StateValue { return []interface{}{} }) + + g := core.NewGraph("fanout-fail").WithStateSchema(schema) + g.AddUpdateNode("good", "good", func(ctx context.Context, s *core.BaseState) (map[string]core.StateValue, error) { + return map[string]core.StateValue{"log": []interface{}{"good"}}, nil + }) + g.AddUpdateNode("bad", "bad", func(ctx context.Context, s *core.BaseState) (map[string]core.StateValue, error) { + return nil, errors.New("branch failed") + }) + g.AddUpdateNode("panicky", "panicky", func(ctx context.Context, s *core.BaseState) (map[string]core.StateValue, error) { + panic("branch exploded") + }) + require.NoError(t, g.SetStartNode("good")) + + out, err := g.ExecuteParallelUpdates(context.Background(), []string{"good", "bad", "panicky"}, schema.NewState()) + require.Error(t, err) + assert.Contains(t, err.Error(), "branch failed") + assert.True(t, errors.Is(err, core.ErrNodePanic), "a panicking branch is reported, not fatal") + + require.NotNil(t, out) + log, _ := out.Get("log") + assert.Equal(t, []interface{}{"good"}, log, "successful branch work is preserved") +} + +// State values must be isolated between the caller and the graph: mutating a +// state after handing it to the engine must not change the run. +func TestConformance_StateIsDeepCopied(t *testing.T) { + in := core.NewBaseState() + nested := map[string]interface{}{"list": []interface{}{1, 2}} + in.Set("nested", nested) + + clone := in.Clone() + nested["list"] = append(nested["list"].([]interface{}), 3) + nested["added"] = true + + got, _ := clone.Get("nested") + gotMap := got.(map[string]interface{}) + assert.Equal(t, []interface{}{1, 2}, gotMap["list"], "clone must not alias caller data") + _, added := gotMap["added"] + assert.False(t, added) +} + +// Values that reflection cannot rebuild (structs with unexported fields such as +// time.Time) must round-trip rather than panic. +func TestConformance_CloneHandlesOpaqueValues(t *testing.T) { + now := time.Now() + in := core.NewBaseState() + in.Set("ts", now) + in.Set("dur", 5*time.Second) + in.Set("nested", map[string]interface{}{"at": now}) + + clone := in.Clone() + ts, ok := clone.Get("ts") + require.True(t, ok) + assert.True(t, now.Equal(ts.(time.Time))) + nested, _ := clone.Get("nested") + assert.True(t, now.Equal(nested.(map[string]interface{})["at"].(time.Time))) +} + +// Self-referential structures must not hang the copier. +func TestConformance_CloneHandlesCycles(t *testing.T) { + cyclic := map[string]interface{}{"name": "root"} + cyclic["self"] = cyclic + + in := core.NewBaseState() + in.Set("cyclic", cyclic) + + done := make(chan struct{}) + go func() { + defer close(done) + _ = in.Clone() + }() + select { + case <-done: + case <-time.After(5 * time.Second): + t.Fatal("Clone did not terminate on a cyclic structure") + } +} + +// State must survive a JSON round-trip; anything less silently empties every +// checkpoint and API response. +func TestConformance_StateJSONRoundTrip(t *testing.T) { + in := core.NewBaseState() + in.Set("counter", 42) + in.Set("messages", []interface{}{msg("1", "hi")}) + in.SetMetadata("thread", "t-1") + + encoded, err := in.ToJSON() + require.NoError(t, err) + assert.Contains(t, string(encoded), "counter") + + out := core.NewBaseState() + require.NoError(t, out.FromJSON(encoded)) + + counter, ok := out.Get("counter") + require.True(t, ok) + assert.EqualValues(t, 42, counter) + thread, ok := out.GetMetadata("thread") + require.True(t, ok) + assert.Equal(t, "t-1", thread) + + // Writing after a round-trip must not panic on a nil map. + out.Set("after", true) +} + +// Concurrent readers and writers of one state must be race-free. +func TestConformance_StateConcurrentAccess(t *testing.T) { + state := core.NewBaseState() + var wg sync.WaitGroup + for i := 0; i < 16; i++ { + wg.Add(1) + go func(i int) { + defer wg.Done() + for j := 0; j < 50; j++ { + state.Set("k", i*j) + _, _ = state.Get("k") + _ = state.Keys() + _ = state.GetAll() + _ = state.Clone() + } + }(i) + } + wg.Wait() +} diff --git a/test/conformance/streaming_subgraph_test.go b/test/conformance/streaming_subgraph_test.go new file mode 100644 index 0000000..dfc6bcd --- /dev/null +++ b/test/conformance/streaming_subgraph_test.go @@ -0,0 +1,350 @@ +// Copyright (c) 2024 GoLangGraph Team +// +// Licensed under the MIT License. See LICENSE file in the project root for full license information. + +package conformance + +import ( + "context" + "errors" + "fmt" + "testing" + "time" + + "github.com/UnicoLab/GoLangGraph/pkg/core" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// LangGraph: streaming yields one update per executed node, in execution order, +// carrying the state after that node. +func TestConformance_StreamingEmitsEveryStepInOrder(t *testing.T) { + g := core.NewGraph("stream") + g.AddNode("a", "A", setNode("a", 1)) + g.AddNode("b", "B", setNode("b", 2)) + g.AddNode("c", "C", setNode("c", 3)) + g.AddEdge("a", "b", nil) + g.AddEdge("b", "c", nil) + require.NoError(t, g.SetStartNode("a")) + require.NoError(t, g.AddEndNode("c")) + + stream := make(chan *core.ExecutionResult, 16) + _, err := g.ExecuteWithOptions(context.Background(), core.NewBaseState(), &core.ExecuteOptions{Stream: stream}) + require.NoError(t, err) + + var seen []string + for result := range stream { + seen = append(seen, result.NodeID) + require.True(t, result.Success) + require.NotNil(t, result.State, "each streamed step carries the state after that node") + v, ok := result.State.Get(result.NodeID) + require.True(t, ok, "step for node %s must include its own write", result.NodeID) + assert.NotNil(t, v) + } + assert.Equal(t, []string{"a", "b", "c"}, seen) +} + +// The per-run stream must close when the run ends, including on failure, so +// consumers are never left waiting. +func TestConformance_StreamClosesOnFailure(t *testing.T) { + g := core.NewGraph("stream-fail") + g.AddNode("a", "A", setNode("a", 1)) + g.AddNode("bad", "Bad", func(ctx context.Context, s *core.BaseState) (*core.BaseState, error) { + return nil, errors.New("nope") + }) + g.AddEdge("a", "bad", nil) + require.NoError(t, g.SetStartNode("a")) + + stream := make(chan *core.ExecutionResult, 16) + _, err := g.ExecuteWithOptions(context.Background(), core.NewBaseState(), &core.ExecuteOptions{Stream: stream}) + require.Error(t, err) + + done := make(chan []*core.ExecutionResult, 1) + go func() { + var got []*core.ExecutionResult + for r := range stream { + got = append(got, r) + } + done <- got + }() + + select { + case got := <-done: + require.Len(t, got, 2) + assert.True(t, got[0].Success) + assert.False(t, got[1].Success, "the failing step must be streamed too") + assert.NotEmpty(t, got[1].ErrorMessage) + case <-time.After(5 * time.Second): + t.Fatal("stream was not closed after a failed run") + } +} + +// Streamed results must be JSON-serialisable, since they cross a WebSocket to +// GoLangGraph Studio. An error that vanishes on the wire is a silent failure. +func TestConformance_StreamResultIsSerializable(t *testing.T) { + g := core.NewGraph("stream-json") + g.AddNode("bad", "Bad", func(ctx context.Context, s *core.BaseState) (*core.BaseState, error) { + return nil, errors.New("provider timeout") + }) + require.NoError(t, g.SetStartNode("bad")) + + stream := make(chan *core.ExecutionResult, 4) + _, err := g.ExecuteWithOptions(context.Background(), core.NewBaseState(), &core.ExecuteOptions{Stream: stream}) + require.Error(t, err) + + result := <-stream + require.NotNil(t, result) + encoded, err := jsonMarshal(result) + require.NoError(t, err) + assert.Contains(t, string(encoded), "provider timeout", + "the failure reason must survive serialisation to clients") + assert.NotContains(t, string(encoded), "goroutine ", "stack traces must not be exposed to clients") +} + +// A slow consumer must not stall graph execution. +func TestConformance_SlowStreamConsumerDoesNotBlockExecution(t *testing.T) { + g := core.NewGraph("stream-slow") + for i := 0; i < 50; i++ { + id := fmt.Sprintf("n%d", i) + g.AddNode(id, id, setNode(id, i)) + if i > 0 { + g.AddEdge(fmt.Sprintf("n%d", i-1), id, nil) + } + } + require.NoError(t, g.SetStartNode("n0")) + require.NoError(t, g.AddEndNode("n49")) + + // A deliberately tiny, unread buffer. + stream := make(chan *core.ExecutionResult, 1) + + done := make(chan error, 1) + go func() { + _, err := g.ExecuteWithOptions(context.Background(), core.NewBaseState(), &core.ExecuteOptions{Stream: stream}) + done <- err + }() + + select { + case err := <-done: + require.NoError(t, err) + case <-time.After(15 * time.Second): + t.Fatal("execution blocked on a slow stream consumer") + } +} + +// LangGraph: a compiled graph can be used as a node of another graph. +func TestConformance_SubgraphAsNode(t *testing.T) { + sub := core.NewGraph("inner") + sub.AddNode("double", "Double", func(ctx context.Context, s *core.BaseState) (*core.BaseState, error) { + v, _ := s.Get("value") + n, _ := v.(int) + s.Set("value", n*2) + s.Set("inner_ran", true) + return s, nil + }) + require.NoError(t, sub.SetStartNode("double")) + require.NoError(t, sub.AddEndNode("double")) + + parent := core.NewGraph("outer") + parent.AddNode("seed", "Seed", setNode("value", 21)) + _, err := parent.AddSubgraph("inner", "Inner", sub, nil) + require.NoError(t, err) + parent.AddNode("report", "Report", setNode("done", true)) + parent.AddEdge("seed", "inner", nil) + parent.AddEdge("inner", "report", nil) + require.NoError(t, parent.SetStartNode("seed")) + require.NoError(t, parent.AddEndNode("report")) + + out, err := parent.Execute(context.Background(), core.NewBaseState()) + require.NoError(t, err) + + value, _ := out.Get("value") + assert.Equal(t, 42, value, "the subgraph result must flow into the parent state") + ran, _ := out.Get("inner_ran") + assert.Equal(t, true, ran) + done, _ := out.Get("done") + assert.Equal(t, true, done) +} + +// Input and output projection lets a subgraph expose a narrow interface. +func TestConformance_SubgraphKeyProjection(t *testing.T) { + sub := core.NewGraph("inner") + sub.AddNode("work", "Work", func(ctx context.Context, s *core.BaseState) (*core.BaseState, error) { + // The parent's secret must not be visible here. + if _, leaked := s.Get("secret"); leaked { + return nil, errors.New("subgraph saw a key outside its declared input") + } + v, _ := s.Get("in") + s.Set("out", fmt.Sprintf("processed:%v", v)) + s.Set("scratch", "internal") + return s, nil + }) + require.NoError(t, sub.SetStartNode("work")) + require.NoError(t, sub.AddEndNode("work")) + + parent := core.NewGraph("outer") + parent.AddNode("seed", "Seed", func(ctx context.Context, s *core.BaseState) (*core.BaseState, error) { + s.Set("in", "payload") + s.Set("secret", "do-not-leak") + return s, nil + }) + _, err := parent.AddSubgraph("inner", "Inner", sub, &core.SubgraphOptions{ + InputKeys: []string{"in"}, + OutputKeys: []string{"out"}, + }) + require.NoError(t, err) + parent.AddEdge("seed", "inner", nil) + require.NoError(t, parent.SetStartNode("seed")) + require.NoError(t, parent.AddEndNode("inner")) + + out, err := parent.Execute(context.Background(), core.NewBaseState()) + require.NoError(t, err) + + result, ok := out.Get("out") + require.True(t, ok) + assert.Equal(t, "processed:payload", result) + _, leaked := out.Get("scratch") + assert.False(t, leaked, "keys outside OutputKeys must not reach the parent") + secret, _ := out.Get("secret") + assert.Equal(t, "do-not-leak", secret, "the parent's own keys survive") +} + +// Namespacing keeps a subgraph's output from colliding with parent keys. +func TestConformance_SubgraphNamespace(t *testing.T) { + sub := core.NewGraph("inner") + sub.AddNode("work", "Work", setNode("result", "inner-value")) + require.NoError(t, sub.SetStartNode("work")) + require.NoError(t, sub.AddEndNode("work")) + + parent := core.NewGraph("outer") + parent.AddNode("seed", "Seed", setNode("result", "parent-value")) + _, err := parent.AddSubgraph("inner", "Inner", sub, &core.SubgraphOptions{Namespace: "inner"}) + require.NoError(t, err) + parent.AddEdge("seed", "inner", nil) + require.NoError(t, parent.SetStartNode("seed")) + require.NoError(t, parent.AddEndNode("inner")) + + out, err := parent.Execute(context.Background(), core.NewBaseState()) + require.NoError(t, err) + + parentValue, _ := out.Get("result") + assert.Equal(t, "parent-value", parentValue, "namespaced output must not clobber parent keys") + nested, ok := out.Get("inner") + require.True(t, ok) + assert.Equal(t, "inner-value", nested.(map[string]core.StateValue)["result"]) +} + +// Subgraph reducers apply when merging back into the parent. +func TestConformance_SubgraphMergeUsesReducers(t *testing.T) { + schema := core.NewStateSchema(). + AddChannel("log", core.Append, func() core.StateValue { return []interface{}{} }) + + sub := core.NewGraph("inner").WithStateSchema(schema) + sub.AddUpdateNode("work", "Work", func(ctx context.Context, s *core.BaseState) (map[string]core.StateValue, error) { + return map[string]core.StateValue{"log": []interface{}{"inner"}}, nil + }) + require.NoError(t, sub.SetStartNode("work")) + require.NoError(t, sub.AddEndNode("work")) + + parent := core.NewGraph("outer").WithStateSchema(schema) + parent.AddUpdateNode("seed", "Seed", func(ctx context.Context, s *core.BaseState) (map[string]core.StateValue, error) { + return map[string]core.StateValue{"log": []interface{}{"parent"}}, nil + }) + _, err := parent.AddSubgraph("inner", "Inner", sub, &core.SubgraphOptions{OutputKeys: []string{"log"}, Schema: schema}) + require.NoError(t, err) + parent.AddEdge("seed", "inner", nil) + require.NoError(t, parent.SetStartNode("seed")) + require.NoError(t, parent.AddEndNode("inner")) + + out, err := parent.Execute(context.Background(), schema.NewState()) + require.NoError(t, err) + + log, _ := out.Get("log") + // The subgraph inherits the parent log, appends to it, and the merge reduces + // the returned slice onto the parent's own copy. + assert.Contains(t, fmt.Sprint(log), "parent") + assert.Contains(t, fmt.Sprint(log), "inner") +} + +// A subgraph failure must identify the subgraph and preserve the cause. +func TestConformance_SubgraphFailurePropagates(t *testing.T) { + sentinel := errors.New("inner exploded") + sub := core.NewGraph("inner") + sub.AddNode("boom", "Boom", func(ctx context.Context, s *core.BaseState) (*core.BaseState, error) { + return nil, sentinel + }) + require.NoError(t, sub.SetStartNode("boom")) + + parent := core.NewGraph("outer") + _, err := parent.AddSubgraph("inner", "Inner", sub, nil) + require.NoError(t, err) + require.NoError(t, parent.SetStartNode("inner")) + + _, err = parent.Execute(context.Background(), core.NewBaseState()) + require.Error(t, err) + assert.ErrorIs(t, err, sentinel) + assert.Contains(t, err.Error(), "inner") +} + +// Recursive graph composition must be rejected at build time rather than +// exhausting the stack at run time. +func TestConformance_SubgraphCycleRejected(t *testing.T) { + a := core.NewGraph("a") + a.AddNode("n", "N", setNode("n", 1)) + require.NoError(t, a.SetStartNode("n")) + require.NoError(t, a.AddEndNode("n")) + + b := core.NewGraph("b") + b.AddNode("m", "M", setNode("m", 1)) + require.NoError(t, b.SetStartNode("m")) + require.NoError(t, b.AddEndNode("m")) + + _, err := a.AddSubgraph("b", "B", b, nil) + require.NoError(t, err) + + // b containing a would close the loop. + _, err = b.AddSubgraph("a", "A", a, nil) + require.Error(t, err) + assert.Contains(t, err.Error(), "cycle") + + // Direct self-nesting is rejected too. + _, err = a.AddSubgraph("self", "Self", a, nil) + require.Error(t, err) +} + +// A subgraph must respect its own recursion limit rather than the parent's. +func TestConformance_SubgraphHasOwnRecursionLimit(t *testing.T) { + sub := core.NewGraph("inner") + sub.Config.MaxIterations = 3 + sub.AddNode("loop", "Loop", func(ctx context.Context, s *core.BaseState) (*core.BaseState, error) { + return s, nil + }) + sub.AddEdge("loop", "loop", nil) + require.NoError(t, sub.SetStartNode("loop")) + + parent := core.NewGraph("outer") + parent.Config.MaxIterations = 100 + _, err := parent.AddSubgraph("inner", "Inner", sub, nil) + require.NoError(t, err) + require.NoError(t, parent.SetStartNode("inner")) + + _, err = parent.Execute(context.Background(), core.NewBaseState()) + require.Error(t, err) + assert.True(t, errors.Is(err, core.ErrRecursionLimit)) +} + +// Topology must include conditional routes and subgraph nodes so visualisers +// and Studio render the reachable graph, not a subset of it. +func TestConformance_TopologyIncludesConditionalRoutes(t *testing.T) { + g := core.NewGraph("topology") + g.AddNode("start", "Start", setNode("start", true)) + g.AddNode("a", "A", setNode("a", true)) + g.AddNode("b", "B", setNode("b", true)) + require.NoError(t, g.SetStartNode("start")) + require.NoError(t, g.AddConditionalEdges("start", + func(ctx context.Context, s *core.BaseState) (string, error) { return "x", nil }, + map[string]string{"x": "a", "y": "b"})) + + topo := g.GetTopology() + assert.ElementsMatch(t, []string{"a", "b"}, topo["start"], + "conditional destinations must appear in the topology") +} diff --git a/test/e2e/docker_test.go b/test/e2e/docker_test.go new file mode 100644 index 0000000..3230736 --- /dev/null +++ b/test/e2e/docker_test.go @@ -0,0 +1,249 @@ +// Copyright (c) 2024 GoLangGraph Team +// +// Licensed under the MIT License. See LICENSE file in the project root for full license information. + +package e2e + +import ( + "bufio" + "os" + "os/exec" + "path/filepath" + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// repoRoot locates the module root from the test's working directory. +func repoRoot(t *testing.T) string { + t.Helper() + dir, err := os.Getwd() + require.NoError(t, err) + for i := 0; i < 6; i++ { + if _, err := os.Stat(filepath.Join(dir, "go.mod")); err == nil { + return dir + } + dir = filepath.Dir(dir) + } + t.Fatal("could not locate the module root") + return "" +} + +// dockerfileInstruction is one parsed line of a Dockerfile. +type dockerfileInstruction struct { + Verb string + Args []string + Line int +} + +// parseDockerfile reads a Dockerfile, joining continuation lines. +func parseDockerfile(t *testing.T, path string) []dockerfileInstruction { + t.Helper() + + file, err := os.Open(path) // #nosec G304 -- path is built from the repository root + require.NoError(t, err) + defer func() { _ = file.Close() }() + + var instructions []dockerfileInstruction + scanner := bufio.NewScanner(file) + scanner.Buffer(make([]byte, 0, 64*1024), 1024*1024) + + var pending string + lineNo, startLine := 0, 0 + for scanner.Scan() { + lineNo++ + line := strings.TrimSpace(scanner.Text()) + if line == "" || strings.HasPrefix(line, "#") { + continue + } + if pending == "" { + startLine = lineNo + } + if strings.HasSuffix(line, `\`) { + pending += strings.TrimSuffix(line, `\`) + " " + continue + } + full := pending + line + pending = "" + + fields := strings.Fields(full) + if len(fields) == 0 { + continue + } + instructions = append(instructions, dockerfileInstruction{ + Verb: strings.ToUpper(fields[0]), + Args: fields[1:], + Line: startLine, + }) + } + require.NoError(t, scanner.Err()) + return instructions +} + +// dockerfiles returns the images built from the repository root context. +func dockerfiles(t *testing.T) []string { + t.Helper() + root := repoRoot(t) + var found []string + for _, name := range []string{"Dockerfile", "Dockerfile.agent"} { + path := filepath.Join(root, name) + if _, err := os.Stat(path); err == nil { + found = append(found, path) + } + } + require.NotEmpty(t, found, "no Dockerfiles found") + return found +} + +// A COPY from the build context must reference a path that exists, or the +// image cannot be built at all. This is checked statically because no +// container runtime is available in every environment the tests run in. +func TestDocker_CopySourcesExist(t *testing.T) { + root := repoRoot(t) + + for _, path := range dockerfiles(t) { + t.Run(filepath.Base(path), func(t *testing.T) { + for _, ins := range parseDockerfile(t, path) { + if ins.Verb != "COPY" && ins.Verb != "ADD" { + continue + } + + args := ins.Args + fromStage := false + filtered := args[:0] + for _, a := range args { + if strings.HasPrefix(a, "--from=") { + fromStage = true + continue + } + if strings.HasPrefix(a, "--") { + continue + } + filtered = append(filtered, a) + } + if fromStage || len(filtered) < 2 { + // Copies from an earlier build stage are not context paths. + continue + } + + for _, src := range filtered[:len(filtered)-1] { + src = strings.Trim(src, `"`) + if strings.ContainsAny(src, "*?[") { + continue // globs are resolved by the builder + } + full := filepath.Join(root, filepath.Clean("/"+src)) + _, err := os.Stat(full) + assert.NoError(t, err, + "%s:%d copies %q from the build context, but it does not exist", + filepath.Base(path), ins.Line, src) + } + } + }) + } +} + +// Images must not run as root. +func TestDocker_RunsAsNonRoot(t *testing.T) { + for _, path := range dockerfiles(t) { + t.Run(filepath.Base(path), func(t *testing.T) { + var lastUser string + for _, ins := range parseDockerfile(t, path) { + if ins.Verb == "USER" && len(ins.Args) > 0 { + lastUser = ins.Args[0] + } + } + require.NotEmpty(t, lastUser, "no USER instruction: the image would run as root") + assert.NotEqual(t, "root", lastUser) + assert.NotEqual(t, "0", lastUser) + }) + } +} + +// The health check must probe the server rather than scanning local +// dependencies: a serving container is healthy even when an optional +// dependency is absent, and a dependency scan would restart it forever. +func TestDocker_HealthcheckProbesTheServer(t *testing.T) { + for _, path := range dockerfiles(t) { + t.Run(filepath.Base(path), func(t *testing.T) { + var healthcheck string + for _, ins := range parseDockerfile(t, path) { + if ins.Verb == "HEALTHCHECK" { + healthcheck = strings.Join(ins.Args, " ") + } + } + require.NotEmpty(t, healthcheck, "image has no HEALTHCHECK") + assert.Contains(t, healthcheck, "--server", + "the health check must probe the serving endpoint") + }) + } +} + +// Every image must define an entrypoint and expose the serving port. +func TestDocker_EntrypointAndPort(t *testing.T) { + for _, path := range dockerfiles(t) { + t.Run(filepath.Base(path), func(t *testing.T) { + verbs := map[string]bool{} + for _, ins := range parseDockerfile(t, path) { + verbs[ins.Verb] = true + } + assert.True(t, verbs["ENTRYPOINT"] || verbs["CMD"], "image has no entrypoint") + assert.True(t, verbs["EXPOSE"], "image does not expose a port") + }) + } +} + +// The binary the image builds and the one it runs must be the same, or the +// container starts and immediately fails. +func TestDocker_BuiltBinaryMatchesEntrypoint(t *testing.T) { + for _, path := range dockerfiles(t) { + t.Run(filepath.Base(path), func(t *testing.T) { + var built, entry string + for _, ins := range parseDockerfile(t, path) { + switch ins.Verb { + case "RUN": + joined := strings.Join(ins.Args, " ") + if idx := strings.Index(joined, " -o "); idx >= 0 { + rest := strings.Fields(joined[idx+4:]) + if len(rest) > 0 { + built = rest[0] + } + } + case "ENTRYPOINT": + entry = strings.Trim(strings.Join(ins.Args, " "), `[]"`) + entry = strings.TrimPrefix(strings.Trim(entry, `"`), "./") + } + } + require.NotEmpty(t, built, "no build output found") + require.NotEmpty(t, entry, "no ENTRYPOINT found") + assert.Equal(t, built, entry, + "the image builds %q but runs %q", built, entry) + }) + } +} + +// When a container runtime is available, actually build the image. Skipped +// otherwise so the suite stays runnable everywhere. +func TestDocker_ImageBuilds(t *testing.T) { + if testing.Short() { + t.Skip("short mode") + } + if _, err := exec.LookPath("docker"); err != nil { + t.Skip("docker binary not available") + } + if err := exec.Command("docker", "info").Run(); err != nil { + t.Skip("docker daemon not available in this environment") + } + + root := repoRoot(t) + for _, path := range dockerfiles(t) { + name := filepath.Base(path) + t.Run(name, func(t *testing.T) { + cmd := exec.Command("docker", "build", "-f", path, "-t", + "golanggraph-test:"+strings.ToLower(strings.ReplaceAll(name, ".", "-")), root) + out, err := cmd.CombinedOutput() + require.NoError(t, err, "docker build failed:\n%s", string(out)) + }) + } +} diff --git a/test/e2e/readme_contract_test.go b/test/e2e/readme_contract_test.go new file mode 100644 index 0000000..d415d82 --- /dev/null +++ b/test/e2e/readme_contract_test.go @@ -0,0 +1,140 @@ +// Copyright (c) 2024 GoLangGraph Team +// +// Licensed under the MIT License. See LICENSE file in the project root for full license information. + +package e2e + +import ( + "context" + "testing" + + "github.com/UnicoLab/GoLangGraph/pkg/agent" + "github.com/UnicoLab/GoLangGraph/pkg/core" + "github.com/UnicoLab/GoLangGraph/pkg/llm" + "github.com/UnicoLab/GoLangGraph/pkg/tools" + "github.com/UnicoLab/GoLangGraph/test/fakes" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// The README's Quick Start is the first code most users run. It was never +// exercised by a test, so nothing would catch it drifting from the API. +// +// These tests follow the documented shape exactly β€” including building +// AgentConfig as a bare literal, which is what the README shows β€” substituting +// only the model provider so they run without a live Ollama. + +func TestReadme_SimpleChatAgent(t *testing.T) { + llmManager := llm.NewProviderManager() + provider := fakes.NewProvider("ollama", "Go is a statically typed language.") + require.NoError(t, llmManager.RegisterProvider("ollama", provider)) + + toolRegistry := tools.NewToolRegistry() + + // Exactly the config shape the README documents: no ID field. + config := &agent.AgentConfig{ + Name: "chat-agent", + Type: agent.AgentTypeChat, + Model: "gemma3:1b", + Provider: "ollama", + SystemPrompt: "You are a helpful AI assistant.", + Temperature: 0.7, + MaxTokens: 500, + } + + chatAgent := agent.NewAgent(config, llmManager, toolRegistry) + require.NotNil(t, chatAgent) + + execution, err := chatAgent.Execute(context.Background(), "Hello! Tell me about Go programming.") + require.NoError(t, err, "the documented quick start must actually run") + + assert.True(t, execution.Success) + assert.Contains(t, execution.Output, "statically typed") + assert.Equal(t, 1, provider.Calls(), "the agent must call the configured provider") + + // An agent built from a literal must still get an identity: AgentManager + // keys by ID, so an empty one means every such agent collides. + assert.NotEmpty(t, chatAgent.GetConfig().ID, + "an agent configured the documented way must still have an ID") +} + +// Two agents built the documented way must be independently addressable. +func TestReadme_LiteralConfiguredAgentsAreDistinct(t *testing.T) { + llmManager := llm.NewProviderManager() + require.NoError(t, llmManager.RegisterProvider("ollama", fakes.NewProvider("ollama", "ok"))) + registry := tools.NewToolRegistry() + + newDocumentedAgent := func(name string) agent.Agent { + return agent.NewAgent(&agent.AgentConfig{ + Name: name, + Type: agent.AgentTypeChat, + Model: "gemma3:1b", + Provider: "ollama", + }, llmManager, registry) + } + + first := newDocumentedAgent("first") + second := newDocumentedAgent("second") + + assert.NotEqual(t, first.GetConfig().ID, second.GetConfig().ID, + "two agents must not share an ID, or one replaces the other when registered") +} + +// The README's graph workflow section documents building a graph directly. +func TestReadme_GraphWorkflow(t *testing.T) { + graph := core.NewGraph("workflow") + + graph.AddNode("start", "Start", func(ctx context.Context, state *core.BaseState) (*core.BaseState, error) { + state.Set("step", "started") + return state, nil + }) + graph.AddNode("process", "Process", func(ctx context.Context, state *core.BaseState) (*core.BaseState, error) { + state.Set("step", "processed") + return state, nil + }) + graph.AddEdge("start", "process", nil) + require.NoError(t, graph.SetStartNode("start")) + require.NoError(t, graph.AddEndNode("process")) + + initial := core.NewBaseState() + initial.Set("input", "hello") + + result, err := graph.Execute(context.Background(), initial) + require.NoError(t, err) + + step, ok := result.Get("step") + require.True(t, ok) + assert.Equal(t, "processed", step) + + input, ok := result.Get("input") + require.True(t, ok, "the caller's initial state must survive execution") + assert.Equal(t, "hello", input) +} + +// A ReAct agent with tools, as the README's second example shows. +func TestReadme_ReActAgentWithTools(t *testing.T) { + llmManager := llm.NewProviderManager() + require.NoError(t, llmManager.RegisterProvider("ollama", + fakes.NewProvider("ollama", "The answer is 4."))) + + registry := tools.NewToolRegistry() + // The README lists these by name; they must exist in a default registry. + for _, name := range []string{"calculator", "web_search", "file_read"} { + _, exists := registry.GetTool(name) + assert.True(t, exists, "the README references the %q tool", name) + } + + config := &agent.AgentConfig{ + Name: "react-agent", + Type: agent.AgentTypeReAct, + Model: "gemma3:1b", + Provider: "ollama", + Tools: []string{"calculator"}, + MaxIterations: 3, + } + + reactAgent := agent.NewAgent(config, llmManager, registry) + execution, err := reactAgent.Execute(context.Background(), "What is 2+2?") + require.NoError(t, err) + assert.NotEmpty(t, execution.ExecutionPath, "a run must record the nodes it visited") +} diff --git a/test/e2e/resilience_test.go b/test/e2e/resilience_test.go new file mode 100644 index 0000000..c0107fb --- /dev/null +++ b/test/e2e/resilience_test.go @@ -0,0 +1,290 @@ +// Copyright (c) 2024 GoLangGraph Team +// +// Licensed under the MIT License. See LICENSE file in the project root for full license information. + +package e2e + +import ( + "context" + "fmt" + "net/http" + "runtime" + "strings" + "sync" + "testing" + "time" + + "github.com/UnicoLab/GoLangGraph/pkg/core" + "github.com/gorilla/websocket" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// settledGoroutines waits for the goroutine count to stabilise, so a background +// worker that is merely slow to exit is not mistaken for a leak. +func settledGoroutines(t *testing.T) int { + t.Helper() + var last int + for i := 0; i < 50; i++ { + runtime.GC() + time.Sleep(50 * time.Millisecond) + current := runtime.NumGoroutine() + if current == last { + return current + } + last = current + } + return last +} + +// Repeated executions must not leak goroutines. A framework that leaks one +// goroutine per run exhausts a long-lived server. +func TestResource_RepeatedExecutionsDoNotLeakGoroutines(t *testing.T) { + live := startServer(t, nil) + + // Warm up so first-use allocations are not counted as leaks. + for i := 0; i < 5; i++ { + status, _ := live.do(t, http.MethodPost, "/api/v1/graphs/studio-workflow/execute", + map[string]interface{}{"input": "warmup"}) + require.Equal(t, http.StatusOK, status) + } + + before := settledGoroutines(t) + + for i := 0; i < 100; i++ { + status, _ := live.do(t, http.MethodPost, "/api/v1/graphs/studio-workflow/execute", + map[string]interface{}{"input": fmt.Sprintf("run-%d", i)}) + require.Equal(t, http.StatusOK, status) + } + + after := settledGoroutines(t) + assert.LessOrEqual(t, after, before+10, + "goroutines grew from %d to %d over 100 executions", before, after) +} + +// WebSocket churn must release its connections and streaming goroutines. +func TestResource_WebSocketChurnDoesNotLeak(t *testing.T) { + live := startServer(t, nil) + wsURL := "ws" + strings.TrimPrefix(live.baseURL, "http") + "/api/v1/ws/graphs/studio-workflow/stream" + + dialRun := func() { + conn, _, err := websocket.DefaultDialer.Dial(wsURL, nil) + require.NoError(t, err) + defer func() { _ = conn.Close() }() + + require.NoError(t, conn.WriteJSON(map[string]interface{}{"type": "execute", "input": "x"})) + require.NoError(t, conn.SetReadDeadline(time.Now().Add(15*time.Second))) + for { + var msg map[string]interface{} + if err := conn.ReadJSON(&msg); err != nil { + return + } + if kind, _ := msg["type"].(string); kind == "complete" || kind == "error" { + return + } + } + } + + for i := 0; i < 5; i++ { + dialRun() + } + before := settledGoroutines(t) + + for i := 0; i < 40; i++ { + dialRun() + } + after := settledGoroutines(t) + + assert.LessOrEqual(t, after, before+10, + "goroutines grew from %d to %d over 40 WebSocket sessions", before, after) +} + +// An interrupted graph must release everything it was holding. +func TestResource_InterruptedRunsAreCleanedUp(t *testing.T) { + live := startServer(t, nil) + + blocking := core.NewGraph("blocking-cleanup") + blocking.Config.EnableStreaming = false + entered := make(chan struct{}, 100) + blocking.AddNode("wait", "Wait", func(ctx context.Context, st *core.BaseState) (*core.BaseState, error) { + entered <- struct{}{} + <-ctx.Done() + return nil, ctx.Err() + }) + require.NoError(t, blocking.SetStartNode("wait")) + live.srv.GraphManager().Register("blocking-cleanup", blocking) + + before := settledGoroutines(t) + + for i := 0; i < 20; i++ { + ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond) + done := make(chan struct{}) + go func() { + defer close(done) + _, _ = blocking.Execute(ctx, core.NewBaseState()) + }() + <-entered + cancel() + <-done + } + + after := settledGoroutines(t) + assert.LessOrEqual(t, after, before+10, + "goroutines grew from %d to %d over 20 canceled runs", before, after) + assert.False(t, blocking.IsRunning(), "no run should still be marked active") +} + +// A long-running graph must complete without unbounded memory growth in its +// history or state. +func TestResource_LongRunningGraph(t *testing.T) { + g := core.NewGraph("long-running") + g.Config.MaxIterations = 5000 + g.Config.EnableStreaming = false + + g.AddNode("step", "Step", func(ctx context.Context, st *core.BaseState) (*core.BaseState, error) { + n, _ := st.Get("n") + count, _ := n.(int) + st.Set("n", count+1) + return st, nil + }) + g.AddNode("done", "Done", func(ctx context.Context, st *core.BaseState) (*core.BaseState, error) { + st.Set("done", true) + return st, nil + }) + require.NoError(t, g.SetStartNode("step")) + require.NoError(t, g.AddEndNode("done")) + require.NoError(t, g.AddConditionalEdges("step", + func(ctx context.Context, st *core.BaseState) (string, error) { + n, _ := st.Get("n") + if count, _ := n.(int); count >= 2000 { + return "exit", nil + } + return "again", nil + }, + map[string]string{"again": "step", "exit": "done"})) + + start := time.Now() + out, err := g.Execute(context.Background(), core.NewBaseState()) + require.NoError(t, err) + + n, _ := out.Get("n") + assert.Equal(t, 2000, n) + done, _ := out.Get("done") + assert.Equal(t, true, done) + t.Logf("2001 node executions in %s", time.Since(start)) + + history := g.GetExecutionHistory() + assert.Len(t, history, 2001, "every step must be recorded exactly once") +} + +// The same request sent many times concurrently must produce the same result +// each time, with no cross-talk between the duplicates. +func TestResource_DuplicateConcurrentRequests(t *testing.T) { + live := startServer(t, nil) + + const copies = 30 + var wg sync.WaitGroup + results := make([]string, copies) + + for i := 0; i < copies; i++ { + wg.Add(1) + go func(i int) { + defer wg.Done() + status, body := live.do(t, http.MethodPost, "/api/v1/graphs/studio-workflow/execute", + map[string]interface{}{"input": "identical request"}) + if status != http.StatusOK { + results[i] = fmt.Sprintf("status %d", status) + return + } + var resp struct { + State map[string]interface{} `json:"state"` + } + decode(t, body, &resp) + results[i] = fmt.Sprint(resp.State["result"]) + }(i) + } + wg.Wait() + + for i, got := range results { + assert.Equal(t, "long", got, "duplicate %d produced a different result", i) + } +} + +// Duplicate agent executions must each get their own execution record. +func TestResource_DuplicateAgentExecutionsAreDistinct(t *testing.T) { + live := startServer(t, nil) + + const copies = 10 + var wg sync.WaitGroup + ids := make([]string, copies) + + for i := 0; i < copies; i++ { + wg.Add(1) + go func(i int) { + defer wg.Done() + status, body := live.do(t, http.MethodPost, "/api/v1/agents/studio-agent/execute", + map[string]string{"input": "same input"}) + if status != http.StatusOK { + return + } + var wrapper struct { + Execution map[string]interface{} `json:"execution"` + } + decode(t, body, &wrapper) + ids[i] = fmt.Sprint(wrapper.Execution["id"]) + }(i) + } + wg.Wait() + + seen := map[string]bool{} + for _, id := range ids { + if id == "" || id == "" { + continue + } + assert.False(t, seen[id], "execution ID %s was reused across concurrent runs", id) + seen[id] = true + } + assert.NotEmpty(t, seen, "at least one execution must have completed") +} + +// The server must survive being started and stopped repeatedly, releasing its +// port each time. +func TestResource_ServerRestartCycles(t *testing.T) { + for i := 0; i < 3; i++ { + live := startServer(t, nil) + + status, _ := live.do(t, http.MethodGet, "/api/v1/health", nil) + require.Equal(t, http.StatusOK, status, "cycle %d", i) + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + err := live.srv.Stop(ctx) + cancel() + require.NoError(t, err, "cycle %d shutdown", i) + } +} + +// A slow provider must not hold a request open past the client's deadline. +func TestResource_SlowProviderDoesNotHangRequest(t *testing.T) { + live := startServer(t, nil) + live.provider.WithDelay(30 * time.Second) + + client := &http.Client{Timeout: 2 * time.Second} + req, err := http.NewRequest(http.MethodPost, + live.baseURL+"/api/v1/agents/studio-agent/execute", + strings.NewReader(`{"input":"hi"}`)) + require.NoError(t, err) + req.Header.Set("Content-Type", "application/json") + + start := time.Now() + resp, err := client.Do(req) + if resp != nil { + _ = resp.Body.Close() + } + elapsed := time.Since(start) + + // Either the request is cut off by the client deadline or the server + // returns; what matters is that it does not run for the provider's delay. + assert.Less(t, elapsed, 10*time.Second, + "request took %s; a slow provider must not pin a connection", elapsed) + _ = err +} diff --git a/test/e2e/studio_compat_test.go b/test/e2e/studio_compat_test.go new file mode 100644 index 0000000..f3d1a88 --- /dev/null +++ b/test/e2e/studio_compat_test.go @@ -0,0 +1,572 @@ +// Copyright (c) 2024 GoLangGraph Team +// +// Licensed under the MIT License. See LICENSE file in the project root for full license information. + +// Package e2e drives a real GoLangGraph server over HTTP and WebSocket. +// +// studio_compat_test.go pins the contract that GoLangGraph Studio depends on. +// Studio is a first-class client: every request it makes is exercised here +// against a live server, and each assertion names the Studio code that reads +// the field, so a server change that would break the console fails here first. +package e2e + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "net" + "net/http" + "strings" + "testing" + "time" + + "github.com/UnicoLab/GoLangGraph/pkg/agent" + "github.com/UnicoLab/GoLangGraph/pkg/core" + "github.com/UnicoLab/GoLangGraph/pkg/llm" + "github.com/UnicoLab/GoLangGraph/pkg/server" + "github.com/UnicoLab/GoLangGraph/pkg/tools" + "github.com/UnicoLab/GoLangGraph/test/fakes" + "github.com/gorilla/websocket" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// liveServer is a running GoLangGraph server plus the pieces under test. +type liveServer struct { + baseURL string + apiKey string + provider *fakes.Provider + agent agent.Agent + srv *server.Server +} + +// startServer boots a real server on a free port, exactly as an operator would. +func startServer(t *testing.T, mutate func(*server.ServerConfig)) *liveServer { + t.Helper() + + listener, err := net.Listen("tcp", "127.0.0.1:0") + require.NoError(t, err) + port := listener.Addr().(*net.TCPAddr).Port + require.NoError(t, listener.Close()) + + cfg := server.DefaultServerConfig() + cfg.Host = "127.0.0.1" + cfg.Port = port + cfg.StaticDir = "" + if mutate != nil { + mutate(cfg) + } + + s := server.NewServer(cfg) + + provider := fakes.NewProvider("fake", "fake response") + providers := llm.NewProviderManager() + require.NoError(t, providers.RegisterProvider("fake", provider)) + s.SetLLMManager(providers) + + // NewToolRegistry already registers the built-in tool set. + registry := tools.NewToolRegistry() + s.SetToolRegistry(registry) + + agents := server.NewAgentManager(providers, registry) + agentCfg := agent.DefaultAgentConfig() + agentCfg.ID = "studio-agent" + agentCfg.Name = "Studio Agent" + agentCfg.Type = agent.AgentTypeChat + agentCfg.Provider = "fake" + agentCfg.Model = "fake-model" + agentCfg.Tools = []string{"calculator"} + instance, err := agents.CreateAgent(agentCfg) + require.NoError(t, err) + s.SetAgentManager(agents) + + // A workflow graph, registered so the graph endpoints have real content. + g := core.NewGraph("studio-workflow") + g.AddNode("ingest", "Ingest", func(ctx context.Context, st *core.BaseState) (*core.BaseState, error) { + in, _ := st.Get("input") + st.Set("ingested", fmt.Sprintf("%v", in)) + return st, nil + }) + g.AddNode("decide", "Decide", func(ctx context.Context, st *core.BaseState) (*core.BaseState, error) { + in, _ := st.Get("ingested") + st.Set("route", map[bool]string{true: "long", false: "short"}[len(fmt.Sprint(in)) > 5]) + return st, nil + }) + g.AddNode("long", "Long path", func(ctx context.Context, st *core.BaseState) (*core.BaseState, error) { + st.Set("result", "long") + return st, nil + }) + g.AddNode("short", "Short path", func(ctx context.Context, st *core.BaseState) (*core.BaseState, error) { + st.Set("result", "short") + return st, nil + }) + g.AddEdge("ingest", "decide", nil) + require.NoError(t, g.AddConditionalEdges("decide", + func(ctx context.Context, st *core.BaseState) (string, error) { + v, _ := st.Get("route") + return fmt.Sprint(v), nil + }, + map[string]string{"long": "long", "short": "short"})) + require.NoError(t, g.SetStartNode("ingest")) + require.NoError(t, g.AddEndNode("long")) + require.NoError(t, g.AddEndNode("short")) + s.GraphManager().Register("studio-workflow", g) + + go func() { + if err := s.Start(); err != nil { + t.Logf("server stopped: %v", err) + } + }() + + live := &liveServer{ + baseURL: fmt.Sprintf("http://127.0.0.1:%d", port), + provider: provider, + agent: instance, + srv: s, + } + if cfg.Security != nil && cfg.Security.RequireAuth && len(cfg.Security.APIKeys) > 0 { + live.apiKey = cfg.Security.APIKeys[0] + } + + live.waitReady(t) + t.Cleanup(func() { + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + _ = s.Stop(ctx) + }) + return live +} + +func (l *liveServer) waitReady(t *testing.T) { + t.Helper() + deadline := time.Now().Add(15 * time.Second) + for time.Now().Before(deadline) { + resp, err := http.Get(l.baseURL + "/api/v1/health") + if err == nil { + _ = resp.Body.Close() + if resp.StatusCode == http.StatusOK { + return + } + } + time.Sleep(20 * time.Millisecond) + } + t.Fatal("server did not become ready") +} + +// do issues a request the way the Studio API client does. +func (l *liveServer) do(t *testing.T, method, path string, body interface{}) (int, []byte) { + t.Helper() + + var reader io.Reader + if body != nil { + raw, err := json.Marshal(body) + require.NoError(t, err) + reader = bytes.NewReader(raw) + } + + req, err := http.NewRequest(method, l.baseURL+path, reader) + require.NoError(t, err) + req.Header.Set("Content-Type", "application/json") + if l.apiKey != "" { + req.Header.Set("X-API-Key", l.apiKey) + } + + resp, err := http.DefaultClient.Do(req) + require.NoError(t, err) + defer func() { _ = resp.Body.Close() }() + + data, err := io.ReadAll(resp.Body) + require.NoError(t, err) + return resp.StatusCode, data +} + +func decode(t *testing.T, data []byte, target interface{}) { + t.Helper() + require.NoError(t, json.Unmarshal(data, target), "response was not valid JSON: %s", string(data)) +} + +// --------------------------------------------------------------------------- +// The endpoints Studio's api/client.ts calls +// --------------------------------------------------------------------------- + +// Studio: client.health() +func TestStudio_Health(t *testing.T) { + live := startServer(t, nil) + + status, body := live.do(t, http.MethodGet, "/api/v1/health", nil) + require.Equal(t, http.StatusOK, status) + + var health struct { + Status string `json:"status"` + Timestamp string `json:"timestamp"` + Version string `json:"version"` + Providers map[string]interface{} `json:"providers"` + } + decode(t, body, &health) + + assert.Equal(t, "healthy", health.Status) + assert.NotEmpty(t, health.Timestamp, "Studio renders the timestamp") + assert.NotEmpty(t, health.Version) +} + +// Studio: client.listAgents() reads res.agents and renders name/type/model, so +// the list must carry configuration objects rather than bare IDs. +func TestStudio_ListAgentsReturnsConfigurations(t *testing.T) { + live := startServer(t, nil) + + status, body := live.do(t, http.MethodGet, "/api/v1/agents", nil) + require.Equal(t, http.StatusOK, status) + + var listed struct { + Agents []map[string]interface{} `json:"agents"` + } + decode(t, body, &listed) + require.Len(t, listed.Agents, 1, "body was: %s", string(body)) + + got := listed.Agents[0] + for _, field := range []string{"id", "name", "type", "model", "provider", "temperature", "max_tokens", "max_iterations", "tools", "enable_streaming", "timeout"} { + assert.Contains(t, got, field, "Studio's AgentConfig requires %q", field) + } + assert.Equal(t, "studio-agent", got["id"]) + assert.Equal(t, "Studio Agent", got["name"]) + assert.Equal(t, "chat", got["type"]) +} + +// Studio: client.getAgent(id) reads res.agent. +func TestStudio_GetAgent(t *testing.T) { + live := startServer(t, nil) + + status, body := live.do(t, http.MethodGet, "/api/v1/agents/studio-agent", nil) + require.Equal(t, http.StatusOK, status) + + var wrapper struct { + Agent map[string]interface{} `json:"agent"` + } + decode(t, body, &wrapper) + assert.Equal(t, "studio-agent", wrapper.Agent["id"]) + + status, body = live.do(t, http.MethodGet, "/api/v1/agents/missing", nil) + assert.Equal(t, http.StatusNotFound, status) + var apiErr struct { + Error string `json:"error"` + } + decode(t, body, &apiErr) + assert.NotEmpty(t, apiErr.Error, "Studio's ApiError reads the error field") +} + +// Studio: client.listTools() reads res.tools as string[]. +func TestStudio_ListTools(t *testing.T) { + live := startServer(t, nil) + + status, body := live.do(t, http.MethodGet, "/api/v1/tools", nil) + require.Equal(t, http.StatusOK, status) + + var listed struct { + Tools []string `json:"tools"` + } + decode(t, body, &listed) + assert.Contains(t, listed.Tools, "calculator") +} + +// Studio: client.listProviders() reads res.providers as ProviderInfo[], with a +// name field. Credentials must never appear. +func TestStudio_ListProvidersReturnsDescriptions(t *testing.T) { + live := startServer(t, nil) + + status, body := live.do(t, http.MethodGet, "/api/v1/providers", nil) + require.Equal(t, http.StatusOK, status) + + var listed struct { + Providers []map[string]interface{} `json:"providers"` + } + decode(t, body, &listed) + require.Len(t, listed.Providers, 1, "body was: %s", string(body)) + + got := listed.Providers[0] + assert.Equal(t, "fake", got["name"], "Studio's ProviderInfo requires name") + assert.Equal(t, "fake-model", got["model"]) + + assert.NotContains(t, got, "api_key", "provider credentials must never be served") + assert.NotContains(t, string(body), "super-secret-key") +} + +// Studio: client.executeAgent(id, input) reads res.execution. +func TestStudio_ExecuteAgent(t *testing.T) { + live := startServer(t, nil) + live.provider.Script("hello from the model") + + status, body := live.do(t, http.MethodPost, "/api/v1/agents/studio-agent/execute", + map[string]string{"input": "hi"}) + require.Equal(t, http.StatusOK, status, "body was: %s", string(body)) + + var wrapper struct { + Execution map[string]interface{} `json:"execution"` + } + decode(t, body, &wrapper) + require.NotEmpty(t, wrapper.Execution, "body was: %s", string(body)) + + // The Go struct tags every field, so the wire format is snake_case like the + // rest of the API. Studio's AgentExecution type mirrors exactly these names. + for _, field := range []string{"id", "timestamp", "input", "output", "duration", "success", "execution_path"} { + assert.Contains(t, wrapper.Execution, field, "Studio's AgentExecution requires %q", field) + } + assert.Equal(t, true, wrapper.Execution["success"]) + assert.Equal(t, "hi", wrapper.Execution["input"]) + assert.Equal(t, "hello from the model", wrapper.Execution["output"]) + + // Studio highlights the nodes that ran from execution_path; an empty list + // leaves its debug view blank for a run that did execute. + path, ok := wrapper.Execution["execution_path"].([]interface{}) + require.True(t, ok, "execution_path must be a list: %s", string(body)) + assert.NotEmpty(t, path, "the nodes that ran must be reported to the debugger") +} + +// Studio: client.getAgentHistory(id) reads res.history. +func TestStudio_AgentHistory(t *testing.T) { + live := startServer(t, nil) + + status, _ := live.do(t, http.MethodPost, "/api/v1/agents/studio-agent/execute", + map[string]string{"input": "remember me"}) + require.Equal(t, http.StatusOK, status) + + status, body := live.do(t, http.MethodGet, "/api/v1/agents/studio-agent/history", nil) + require.Equal(t, http.StatusOK, status) + + var wrapper struct { + History []map[string]interface{} `json:"history"` + } + decode(t, body, &wrapper) + require.NotEmpty(t, wrapper.History, "an executed agent must have history: %s", string(body)) + assert.Equal(t, "remember me", wrapper.History[0]["input"]) +} + +// Studio: client.getGraphTopology(id) reads res.topology.nodes / .edges, and +// maps n.id, n.name, n.type, e.from and e.to. +func TestStudio_GraphTopologyShape(t *testing.T) { + live := startServer(t, nil) + + status, body := live.do(t, http.MethodGet, "/api/v1/graphs/studio-workflow/topology", nil) + require.Equal(t, http.StatusOK, status) + + var resp struct { + GraphID string `json:"graph_id"` + Topology struct { + Nodes []map[string]interface{} `json:"nodes"` + Edges []map[string]interface{} `json:"edges"` + } `json:"topology"` + } + decode(t, body, &resp) + + assert.Equal(t, "studio-workflow", resp.GraphID) + require.NotEmpty(t, resp.Topology.Nodes, "Studio falls back to a synthetic graph when nodes are empty") + require.NotEmpty(t, resp.Topology.Edges) + + for _, n := range resp.Topology.Nodes { + assert.Contains(t, n, "id") + assert.Contains(t, n, "name") + assert.Contains(t, n, "type") + } + for _, e := range resp.Topology.Edges { + assert.Contains(t, e, "from") + assert.Contains(t, e, "to") + } + + // Conditional destinations must be present, or the rendered graph is wrong. + var sawLong, sawShort bool + for _, e := range resp.Topology.Edges { + if e["from"] == "decide" && e["to"] == "long" { + sawLong = true + } + if e["from"] == "decide" && e["to"] == "short" { + sawShort = true + } + } + assert.True(t, sawLong && sawShort, "both conditional routes must appear: %v", resp.Topology.Edges) +} + +// Studio requests a topology using the *agent* ID, so an agent's execution +// graph must resolve through the same endpoint. +func TestStudio_GraphTopologyResolvesAgentID(t *testing.T) { + live := startServer(t, nil) + + status, body := live.do(t, http.MethodGet, "/api/v1/graphs/studio-agent/topology", nil) + require.Equal(t, http.StatusOK, status, "body was: %s", string(body)) + + var resp struct { + Topology struct { + Nodes []map[string]interface{} `json:"nodes"` + } `json:"topology"` + } + decode(t, body, &resp) + assert.NotEmpty(t, resp.Topology.Nodes, + "an agent's own graph must be reachable by its ID, or Studio's graph view is always empty") +} + +// --------------------------------------------------------------------------- +// Cross-origin and authentication, as a browser-hosted Studio experiences them +// --------------------------------------------------------------------------- + +func TestStudio_CORSPreflightAndRequest(t *testing.T) { + origin := "http://localhost:3000" + live := startServer(t, func(c *server.ServerConfig) { + c.Security.AllowedOrigins = []string{origin} + }) + + // Preflight for the JSON POST Studio makes. + req, err := http.NewRequest(http.MethodOptions, live.baseURL+"/api/v1/agents/studio-agent/execute", nil) + require.NoError(t, err) + req.Header.Set("Origin", origin) + req.Header.Set("Access-Control-Request-Method", "POST") + req.Header.Set("Access-Control-Request-Headers", "content-type,x-api-key") + + resp, err := http.DefaultClient.Do(req) + require.NoError(t, err) + defer func() { _ = resp.Body.Close() }() + + require.Less(t, resp.StatusCode, 300, "preflight must succeed or the browser blocks every call") + assert.Equal(t, origin, resp.Header.Get("Access-Control-Allow-Origin")) + assert.Contains(t, resp.Header.Get("Access-Control-Allow-Headers"), "X-API-Key") + assert.Contains(t, resp.Header.Get("Access-Control-Allow-Methods"), "POST") +} + +func TestStudio_AuthenticatedSession(t *testing.T) { + live := startServer(t, func(c *server.ServerConfig) { + c.Security.RequireAuth = true + c.Security.APIKeys = []string{"studio-key"} + }) + + // With the key, Studio works. + status, _ := live.do(t, http.MethodGet, "/api/v1/agents", nil) + assert.Equal(t, http.StatusOK, status) + + // Without it, the server refuses. + req, err := http.NewRequest(http.MethodGet, live.baseURL+"/api/v1/agents", nil) + require.NoError(t, err) + resp, err := http.DefaultClient.Do(req) + require.NoError(t, err) + defer func() { _ = resp.Body.Close() }() + assert.Equal(t, http.StatusUnauthorized, resp.StatusCode) + + // Health stays reachable so the connection screen can probe the server. + probe, err := http.Get(live.baseURL + "/api/v1/health") + require.NoError(t, err) + defer func() { _ = probe.Body.Close() }() + assert.Equal(t, http.StatusOK, probe.StatusCode) +} + +// --------------------------------------------------------------------------- +// Debugging a workflow, which is what Studio exists to do +// --------------------------------------------------------------------------- + +func TestStudio_ExecuteGraphReturnsPerStepDebugInfo(t *testing.T) { + live := startServer(t, nil) + + status, body := live.do(t, http.MethodPost, "/api/v1/graphs/studio-workflow/execute", + map[string]interface{}{"input": "a longer message"}) + require.Equal(t, http.StatusOK, status, "body was: %s", string(body)) + + var resp struct { + Status string `json:"status"` + State map[string]interface{} `json:"state"` + Steps []struct { + NodeID string `json:"node_id"` + Step int `json:"step"` + Success bool `json:"success"` + State map[string]interface{} `json:"state"` + } `json:"steps"` + } + decode(t, body, &resp) + + assert.Equal(t, "completed", resp.Status) + assert.Equal(t, "long", resp.State["result"], "conditional routing must pick the long path") + + require.Len(t, resp.Steps, 3, "a debugger needs every node visit: %v", resp.Steps) + assert.Equal(t, []string{"ingest", "decide", "long"}, + []string{resp.Steps[0].NodeID, resp.Steps[1].NodeID, resp.Steps[2].NodeID}) + + for i, step := range resp.Steps { + assert.True(t, step.Success) + assert.Equal(t, i, step.Step) + assert.NotEmpty(t, step.State, "each step must carry the state after it, for step-through debugging") + } +} + +// A live graph run streamed over WebSocket, the way an interactive debugger +// would drive it. +func TestStudio_WebSocketGraphRun(t *testing.T) { + live := startServer(t, nil) + + wsURL := "ws" + strings.TrimPrefix(live.baseURL, "http") + "/api/v1/ws/graphs/studio-workflow/stream" + conn, resp, err := websocket.DefaultDialer.Dial(wsURL, nil) + if err != nil && resp != nil { + t.Fatalf("dial failed: %v (status %s)", err, resp.Status) + } + require.NoError(t, err) + defer func() { _ = conn.Close() }() + + require.NoError(t, conn.WriteJSON(map[string]interface{}{ + "type": "execute", "input": "tiny", + })) + require.NoError(t, conn.SetReadDeadline(time.Now().Add(20*time.Second))) + + var kinds, nodes []string + var finalState map[string]interface{} + for { + var msg map[string]interface{} + require.NoError(t, conn.ReadJSON(&msg)) + kind, _ := msg["type"].(string) + kinds = append(kinds, kind) + + if kind == "step" { + step := msg["step"].(map[string]interface{}) + nodes = append(nodes, step["node_id"].(string)) + } + if kind == "complete" { + finalState, _ = msg["state"].(map[string]interface{}) + break + } + if kind == "error" { + t.Fatalf("unexpected error frame: %v", msg) + } + } + + assert.Equal(t, "start", kinds[0]) + assert.Equal(t, []string{"ingest", "decide", "short"}, nodes, + "a short input must take the short branch") + require.NotNil(t, finalState) + assert.Equal(t, "short", finalState["result"]) +} + +// A failing provider must surface as a readable error, not a hang or a 500 with +// no explanation. +func TestStudio_ProviderFailureIsReported(t *testing.T) { + live := startServer(t, nil) + live.provider.FailWith(fmt.Errorf("model backend is offline")) + + status, body := live.do(t, http.MethodPost, "/api/v1/agents/studio-agent/execute", + map[string]string{"input": "hi"}) + + assert.GreaterOrEqual(t, status, 400, "a provider failure must not report success") + assert.Contains(t, strings.ToLower(string(body)), "offline", + "Studio shows the server's error field to the user: %s", string(body)) + assert.NotContains(t, string(body), "goroutine ", "stack traces must not reach the console") +} + +// A failed execution must carry a readable reason. A Go error field marshals to +// an empty object, which would leave Studio showing a failure with no cause. +func TestStudio_FailedExecutionCarriesReason(t *testing.T) { + live := startServer(t, nil) + live.provider.FailWith(fmt.Errorf("context length exceeded")) + + _, body := live.do(t, http.MethodPost, "/api/v1/agents/studio-agent/execute", + map[string]string{"input": "hi"}) + + // Whether the failure is reported on the execution object or as a top-level + // error, the reason itself must survive serialisation. + assert.Contains(t, string(body), "context length exceeded", + "the failure reason must reach the client: %s", string(body)) + assert.NotContains(t, string(body), `"error":{}`, + "a Go error must not serialize as an empty object") +} diff --git a/test/fakes/provider.go b/test/fakes/provider.go new file mode 100644 index 0000000..ebda04e --- /dev/null +++ b/test/fakes/provider.go @@ -0,0 +1,197 @@ +// Copyright (c) 2024 GoLangGraph Team +// +// Licensed under the MIT License. See LICENSE file in the project root for full license information. + +// Package fakes provides deterministic test doubles for GoLangGraph's external +// dependencies. +// +// The doubles stand in for a language model only. Everything else in a test +// using them β€” the graph engine, the agent loop, tool execution, state +// handling, the HTTP server β€” is the real implementation, so the tests +// exercise the framework rather than a mock of it. +package fakes + +import ( + "context" + "fmt" + "sync" + "sync/atomic" + "time" + + "github.com/UnicoLab/GoLangGraph/pkg/llm" +) + +// Provider is a scriptable llm.Provider. +// +// Replies are returned in order and the last one repeats, so a test can script +// an agent's reasoning turn by turn without depending on a live model. +type Provider struct { + name string + + mu sync.Mutex + replies []string + failures []error + delay time.Duration + prompts []string + + calls atomic.Int32 + healthy atomic.Bool +} + +// NewProvider creates a provider that always answers with reply. +func NewProvider(name, reply string) *Provider { + p := &Provider{name: name, replies: []string{reply}} + p.healthy.Store(true) + return p +} + +// Script sets the replies returned by successive calls. The final reply repeats +// once the script is exhausted. +func (p *Provider) Script(replies ...string) *Provider { + p.mu.Lock() + defer p.mu.Unlock() + p.replies = append([]string(nil), replies...) + return p +} + +// FailWith makes the next calls fail with the given errors, in order. A nil +// entry means that call succeeds. +func (p *Provider) FailWith(errs ...error) *Provider { + p.mu.Lock() + defer p.mu.Unlock() + p.failures = append([]error(nil), errs...) + return p +} + +// WithDelay makes every call take at least d, respecting cancellation. +func (p *Provider) WithDelay(d time.Duration) *Provider { + p.mu.Lock() + defer p.mu.Unlock() + p.delay = d + return p +} + +// SetHealthy controls what IsHealthy reports. +func (p *Provider) SetHealthy(healthy bool) *Provider { + p.healthy.Store(healthy) + return p +} + +// Calls returns how many completions have been requested. +func (p *Provider) Calls() int { return int(p.calls.Load()) } + +// Prompts returns the content of every user message the provider has received, +// so a test can assert what the agent actually asked. +func (p *Provider) Prompts() []string { + p.mu.Lock() + defer p.mu.Unlock() + return append([]string(nil), p.prompts...) +} + +// next returns the scripted reply and failure for a call index. +func (p *Provider) next(index int, prompt string) (string, error, time.Duration) { + p.mu.Lock() + defer p.mu.Unlock() + + p.prompts = append(p.prompts, prompt) + + var failure error + if index < len(p.failures) { + failure = p.failures[index] + } + + reply := "" + switch { + case len(p.replies) == 0: + case index < len(p.replies): + reply = p.replies[index] + default: + reply = p.replies[len(p.replies)-1] + } + return reply, failure, p.delay +} + +func (p *Provider) GetName() string { return p.name } + +func (p *Provider) GetModels(ctx context.Context) ([]string, error) { + return []string{"fake-model"}, nil +} + +func (p *Provider) Complete(ctx context.Context, req llm.CompletionRequest) (*llm.CompletionResponse, error) { + index := int(p.calls.Add(1)) - 1 + + prompt := "" + for _, m := range req.Messages { + if m.Role == "user" { + prompt = m.Content + } + } + + reply, failure, delay := p.next(index, prompt) + + if delay > 0 { + select { + case <-ctx.Done(): + return nil, ctx.Err() + case <-time.After(delay): + } + } + if err := ctx.Err(); err != nil { + return nil, err + } + if failure != nil { + return nil, failure + } + + return &llm.CompletionResponse{ + ID: fmt.Sprintf("%s-%d", p.name, index), + Object: "chat.completion", + Created: time.Now().Unix(), + Model: req.Model, + Choices: []llm.Choice{{ + Index: 0, + Message: llm.Message{Role: "assistant", Content: reply}, + FinishReason: "stop", + }}, + Usage: llm.Usage{PromptTokens: 1, CompletionTokens: 1, TotalTokens: 2}, + }, nil +} + +func (p *Provider) CompleteStream(ctx context.Context, req llm.CompletionRequest, callback llm.StreamCallback) error { + resp, err := p.Complete(ctx, req) + if err != nil { + return err + } + return callback(*resp) +} + +func (p *Provider) CompleteWithMode(ctx context.Context, req llm.CompletionRequest, mode llm.StreamMode) (*llm.CompletionResponse, error) { + return p.Complete(ctx, req) +} + +func (p *Provider) CompleteStreamWithMode(ctx context.Context, req llm.CompletionRequest, callback llm.StreamCallback, mode llm.StreamMode) error { + return p.CompleteStream(ctx, req, callback) +} + +func (p *Provider) IsHealthy(ctx context.Context) error { + if !p.healthy.Load() { + return fmt.Errorf("provider %s is unhealthy", p.name) + } + return nil +} + +func (p *Provider) GetConfig() map[string]interface{} { + return map[string]interface{}{ + "type": "fake", + "endpoint": "memory://fake", + "model": "fake-model", + // Present on purpose: the API must never echo credentials. + "api_key": "super-secret-key", + } +} + +func (p *Provider) SetConfig(config map[string]interface{}) error { return nil } +func (p *Provider) SupportsStreaming() bool { return true } +func (p *Provider) GetStreamingConfig() *llm.StreamingConfig { return llm.DefaultStreamingConfig() } +func (p *Provider) SetStreamingConfig(c *llm.StreamingConfig) error { return nil } +func (p *Provider) Close() error { return nil }