diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 2d79d73..910c905 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -229,16 +229,24 @@ jobs: working-directory: webui - name: Build Linux - run: GOOS=linux GOARCH=amd64 go build -ldflags "-s -w" -o xsql-linux-amd64 ./cmd/xsql + run: | + GOOS=linux GOARCH=amd64 go build -ldflags "-s -w" -o xsql-linux-amd64 ./cmd/xsql + GOOS=linux GOARCH=amd64 go build -ldflags "-s -w" -o xsql-ai-linux-amd64 ./cmd/xsql-ai - name: Build Windows - run: GOOS=windows GOARCH=amd64 go build -ldflags "-s -w" -o xsql-windows-amd64.exe ./cmd/xsql + run: | + GOOS=windows GOARCH=amd64 go build -ldflags "-s -w" -o xsql-windows-amd64.exe ./cmd/xsql + GOOS=windows GOARCH=amd64 go build -ldflags "-s -w" -o xsql-ai-windows-amd64.exe ./cmd/xsql-ai - name: Build macOS (amd64) - run: GOOS=darwin GOARCH=amd64 go build -ldflags "-s -w" -o xsql-darwin-amd64 ./cmd/xsql + run: | + GOOS=darwin GOARCH=amd64 go build -ldflags "-s -w" -o xsql-darwin-amd64 ./cmd/xsql + GOOS=darwin GOARCH=amd64 go build -ldflags "-s -w" -o xsql-ai-darwin-amd64 ./cmd/xsql-ai - name: Build macOS (arm64) - run: GOOS=darwin GOARCH=arm64 go build -ldflags "-s -w" -o xsql-darwin-arm64 ./cmd/xsql + run: | + GOOS=darwin GOARCH=arm64 go build -ldflags "-s -w" -o xsql-darwin-arm64 ./cmd/xsql + GOOS=darwin GOARCH=arm64 go build -ldflags "-s -w" -o xsql-ai-darwin-arm64 ./cmd/xsql-ai - name: Upload artifacts uses: actions/upload-artifact@ea165f8d65b6e75b540449e92b4886f43607fa02 # v4.6.2 diff --git a/.gitignore b/.gitignore index 0b6494c..373e214 100644 --- a/.gitignore +++ b/.gitignore @@ -11,3 +11,4 @@ coverage*.txt coverage*.html coverage.out stats_coverage.out +bin/ diff --git a/.golangci.yml b/.golangci.yml index 6fe9a3c..c4e5ff7 100644 --- a/.golangci.yml +++ b/.golangci.yml @@ -1,5 +1,5 @@ -version: "2" - +version: "2" + run: timeout: 5m tests: false @@ -7,17 +7,20 @@ run: linters: default: standard settings: - errcheck: - exclude-functions: - - (*database/sql.Rows).Close - - (*database/sql.DB).Close - - (*github.com/zx06/xsql/internal/ssh.Client).Close - - (io.Closer).Close - -formatters: - enable: - - goimports - settings: - goimports: - local-prefixes: - - github.com/zx06/xsql + errcheck: + exclude-functions: + - (*database/sql.Rows).Close + - (*database/sql.DB).Close + - (*github.com/zx06/xsql/internal/ssh.Client).Close + - (io.Closer).Close + - (*os.File).Close + - (*net/http.Response.Body).Close + - io.Copy + +formatters: + enable: + - goimports + settings: + goimports: + local-prefixes: + - github.com/zx06/xsql diff --git a/.goreleaser.yaml b/.goreleaser.yaml index 8dac074..8c7a3a6 100644 --- a/.goreleaser.yaml +++ b/.goreleaser.yaml @@ -1,86 +1,105 @@ -version: 2 - -project_name: xsql - +version: 2 + +project_name: xsql + before: hooks: - npm --prefix webui ci - npm --prefix webui run build - go mod tidy - -builds: - - id: xsql - main: ./cmd/xsql - binary: xsql - env: - - CGO_ENABLED=0 - goos: - - linux - - darwin - - windows - goarch: - - amd64 - - arm64 - ldflags: - - -s -w - - -X main.version={{.Version}} - - -X main.commit={{.ShortCommit}} - - -X main.date={{.Date}} - -archives: - - id: default - name_template: "{{ .ProjectName }}_{{ .Version }}_{{ .Os }}_{{ .Arch }}" - format_overrides: - - goos: windows - format: zip - files: - - README.md - - LICENSE* - - docs/* - -checksum: - name_template: "checksums.txt" - -snapshot: - version_template: "{{ incpatch .Version }}-next" - -changelog: - sort: asc - filters: - exclude: - - "^docs:" - - "^test:" - - "^chore:" - - "Merge pull request" - - "Merge branch" - -release: - github: - owner: zx06 - name: xsql - draft: false - prerelease: auto - name_template: "v{{.Version}}" - -brews: - - repository: - owner: zx06 - name: homebrew-tap - token: "{{ .Env.HOMEBREW_TAP_TOKEN }}" - directory: Formula - homepage: "https://github.com/zx06/xsql" - description: "AI-first cross-database CLI tool with SSH proxy support" - license: "MIT" - install: | - bin.install "xsql" - test: | - system "#{bin}/xsql", "version" - -scoops: - - repository: - owner: zx06 - name: scoop-bucket - token: "{{ .Env.SCOOP_BUCKET_TOKEN }}" - homepage: "https://github.com/zx06/xsql" - description: "AI-first cross-database CLI tool with SSH proxy support" - license: "MIT" + +builds: + - id: xsql + main: ./cmd/xsql + binary: xsql + env: + - CGO_ENABLED=0 + goos: + - linux + - darwin + - windows + goarch: + - amd64 + - arm64 + ldflags: + - -s -w + - -X main.version={{.Version}} + - -X main.commit={{.ShortCommit}} + - -X main.date={{.Date}} + + - id: xsql-ai + main: ./cmd/xsql-ai + binary: xsql-ai + env: + - CGO_ENABLED=0 + goos: + - linux + - darwin + - windows + goarch: + - amd64 + - arm64 + ldflags: + - -s -w + - -X main.version={{.Version}} + - -X main.commit={{.ShortCommit}} + - -X main.date={{.Date}} + +archives: + - id: default + name_template: "{{ .ProjectName }}_{{ .Version }}_{{ .Os }}_{{ .Arch }}" + format_overrides: + - goos: windows + format: zip + files: + - README.md + - LICENSE* + - docs/* + +checksum: + name_template: "checksums.txt" + +snapshot: + version_template: "{{ incpatch .Version }}-next" + +changelog: + sort: asc + filters: + exclude: + - "^docs:" + - "^test:" + - "^chore:" + - "Merge pull request" + - "Merge branch" + +release: + github: + owner: zx06 + name: xsql + draft: false + prerelease: auto + name_template: "v{{.Version}}" + +brews: + - repository: + owner: zx06 + name: homebrew-tap + token: "{{ .Env.HOMEBREW_TAP_TOKEN }}" + directory: Formula + homepage: "https://github.com/zx06/xsql" + description: "AI-first cross-database CLI tool with SSH proxy support" + license: "MIT" + install: | + bin.install "xsql" + bin.install "xsql-ai" + test: | + system "#{bin}/xsql", "version" + +scoops: + - repository: + owner: zx06 + name: scoop-bucket + token: "{{ .Env.SCOOP_BUCKET_TOKEN }}" + homepage: "https://github.com/zx06/xsql" + description: "AI-first cross-database CLI tool with SSH proxy support" + license: "MIT" diff --git a/cmd/xsql-ai/main.go b/cmd/xsql-ai/main.go new file mode 100644 index 0000000..c55d98c --- /dev/null +++ b/cmd/xsql-ai/main.go @@ -0,0 +1,107 @@ +package main + +import ( + "fmt" + "os" + "strings" + + tea "github.com/charmbracelet/bubbletea" + "github.com/spf13/cobra" + + "github.com/zx06/xsql/internal/ai" + "github.com/zx06/xsql/internal/config" + _ "github.com/zx06/xsql/internal/db/mysql" + _ "github.com/zx06/xsql/internal/db/pg" + "github.com/zx06/xsql/internal/secret" + "github.com/zx06/xsql/internal/tui" +) + +type AIFlags struct { + ConfigPath string + Profile string + Model string + BaseURL string + APIKey string + UnsafeAllowWrite bool + Prompt string +} + +var newProgramFunc = func(model tea.Model) *tea.Program { + return tea.NewProgram(model, tea.WithAltScreen()) +} + +func newRootCmd() *cobra.Command { + flags := &AIFlags{} + + rootCmd := &cobra.Command{ + Use: "xsql-ai [PROMPT]", + Short: "xsql AI interactive database query tool (TUI)", + RunE: func(cmd *cobra.Command, args []string) error { + if len(args) > 0 { + flags.Prompt = strings.Join(args, " ") + } + return runAI(cmd, flags) + }, + } + + rootCmd.Flags().StringVar(&flags.ConfigPath, "config", "", "Config file path (YAML)") + rootCmd.Flags().StringVarP(&flags.Profile, "profile", "p", "", "Profile name (default: 'default')") + rootCmd.Flags().StringVar(&flags.Model, "model", "", "AI model name (default: gpt-4o)") + rootCmd.Flags().StringVar(&flags.BaseURL, "base-url", "", "AI service base URL") + rootCmd.Flags().StringVar(&flags.APIKey, "api-key", "", "AI service API key") + rootCmd.Flags().BoolVar(&flags.UnsafeAllowWrite, "unsafe-allow-write", false, "Allow write operations (bypasses read-only protection)") + rootCmd.Flags().StringVar(&flags.Prompt, "prompt", "", "Initial prompt for AI query") + + return rootCmd +} + +func main() { + rootCmd := newRootCmd() + if err := rootCmd.Execute(); err != nil { + os.Exit(1) + } +} + +func runAI(cmd *cobra.Command, flags *AIFlags) error { + opts := config.Options{ + ConfigPath: flags.ConfigPath, + CLIProfile: flags.Profile, + CLIProfileSet: flags.Profile != "", + CLIAIModel: flags.Model, + CLIAIModelSet: cmd.Flags().Changed("model"), + CLIAIBaseURL: flags.BaseURL, + CLIAIBaseURLSet: cmd.Flags().Changed("base-url"), + CLIAIAPIKey: flags.APIKey, + CLIAIAPIKeySet: cmd.Flags().Changed("api-key"), + } + + resolved, xe := config.Resolve(opts) + if xe != nil { + return fmt.Errorf("config error [%s]: %s", xe.Code, xe.Message) + } + if resolved.ProfileName == "" || resolved.Profile.DB == "" { + return fmt.Errorf("config error [XSQL_CFG_INVALID]: no profile specified and no 'default' profile found in config") + } + + // Resolve API key if keyring reference or plaintext + apiKey := resolved.AI.APIKey + if secret.IsKeyringRef(apiKey) { + resolvedKey, xe := secret.Resolve(apiKey, secret.Options{AllowPlaintext: true}) + if xe == nil { + apiKey = resolvedKey + } + } + resolved.AI.APIKey = apiKey + + aiClient := ai.NewClient(resolved.AI, nil) + aiService := ai.NewService(resolved.AI, aiClient) + + model := tui.NewModel(opts, resolved, aiService, flags.Prompt, flags.UnsafeAllowWrite) + + p := newProgramFunc(model) + if _, err := p.Run(); err != nil { + return fmt.Errorf("error running TUI: %w", err) + } + + return nil +} diff --git a/cmd/xsql-ai/main_test.go b/cmd/xsql-ai/main_test.go new file mode 100644 index 0000000..39e29ed --- /dev/null +++ b/cmd/xsql-ai/main_test.go @@ -0,0 +1,138 @@ +package main + +import ( + "os" + "path/filepath" + "strings" + "testing" + + tea "github.com/charmbracelet/bubbletea" + + "github.com/zx06/xsql/internal/config" +) + +func TestCmdXSQLAI_MissingProfileError(t *testing.T) { + tmpDir := t.TempDir() + emptyConfig := filepath.Join(tmpDir, "empty.yaml") + if err := os.WriteFile(emptyConfig, []byte("profiles: {}\n"), 0600); err != nil { + t.Fatalf("failed to write empty config: %v", err) + } + + cmd := newRootCmd() + cmd.SetArgs([]string{"--config", emptyConfig}) + + err := cmd.Execute() + if err == nil { + t.Fatal("expected error when no profile specified and no default profile") + } + if !strings.Contains(err.Error(), "no profile specified") { + t.Fatalf("expected missing profile error message, got %v", err) + } +} + +func TestCmdXSQLAI_ConfigPathNotFound(t *testing.T) { + cmd := newRootCmd() + cmd.SetArgs([]string{"--config", "/nonexistent/config.yaml"}) + + err := cmd.Execute() + if err == nil { + t.Fatal("expected error when config file does not exist") + } + if !strings.Contains(err.Error(), "config error") { + t.Fatalf("expected config error prefix, got %v", err) + } +} + +func TestCmdXSQLAI_FlagsBinding(t *testing.T) { + tmpDir := t.TempDir() + cfgPath := filepath.Join(tmpDir, "xsql.yaml") + cfgContent := ` +profiles: + dev: + db: mysql + host: 127.0.0.1 + port: 3306 + user: root + database: testdb +` + if err := os.WriteFile(cfgPath, []byte(cfgContent), 0600); err != nil { + t.Fatalf("failed to write temp config: %v", err) + } + + flags := &AIFlags{ + ConfigPath: cfgPath, + Profile: "dev", + Model: "gpt-4o", + } + + cmd := newRootCmd() + cmd.SetArgs([]string{"--config", cfgPath, "--profile", "dev", "Show top 10 servers"}) + + // Verify command flag parsing + if err := cmd.ParseFlags([]string{"--config", cfgPath, "--profile", "dev", "--unsafe-allow-write"}); err != nil { + t.Fatalf("failed to parse flags: %v", err) + } + + if flags.Profile != "dev" { + t.Fatalf("expected profile 'dev', got %q", flags.Profile) + } +} + +func TestCmdXSQLAI_RunAIValidation(t *testing.T) { + tmpDir := t.TempDir() + cfgPath := filepath.Join(tmpDir, "xsql.yaml") + cfgContent := ` +ai: + api_key: "keyring:test_key" +profiles: + dev: + db: mysql + host: 127.0.0.1 + port: 3306 + user: root + database: testdb +` + if err := os.WriteFile(cfgPath, []byte(cfgContent), 0600); err != nil { + t.Fatalf("failed to write temp config: %v", err) + } + + flags := &AIFlags{ + ConfigPath: cfgPath, + Profile: "dev", + APIKey: "test-key", + } + + cmd := newRootCmd() + cmd.SetArgs([]string{"--config", cfgPath, "--profile", "dev", "--api-key", "test-key", "show users"}) + _ = cmd.ParseFlags([]string{"--config", cfgPath, "--profile", "dev", "--api-key", "test-key"}) + + // Validate runAI options resolution up to TUI program + opts := config.Options{ + ConfigPath: flags.ConfigPath, + CLIProfile: flags.Profile, + CLIProfileSet: true, + CLIAIAPIKey: flags.APIKey, + } + resolved, xe := config.Resolve(opts) + if xe != nil { + t.Fatalf("unexpected error resolving config options: %v", xe) + } + if resolved.ProfileName != "dev" { + t.Fatalf("expected profile 'dev', got %q", resolved.ProfileName) + } + + // Test full runAI execution flow + oldNewProgram := newProgramFunc + defer func() { newProgramFunc = oldNewProgram }() + + newProgramFunc = func(model tea.Model) *tea.Program { + p := tea.NewProgram(model, tea.WithInput(strings.NewReader("")), tea.WithOutput(os.Stderr), tea.WithoutRenderer()) + go p.Quit() + return p + } + + err := cmd.Execute() + if err != nil { + t.Fatalf("expected runAI to execute successfully, got %v", err) + } +} diff --git a/cmd/xsql/ai.go b/cmd/xsql/ai.go new file mode 100644 index 0000000..8b8963f --- /dev/null +++ b/cmd/xsql/ai.go @@ -0,0 +1,83 @@ +package main + +import ( + "fmt" + "strings" + + tea "github.com/charmbracelet/bubbletea" + "github.com/spf13/cobra" + + "github.com/zx06/xsql/internal/ai" + "github.com/zx06/xsql/internal/config" + "github.com/zx06/xsql/internal/secret" + "github.com/zx06/xsql/internal/tui" +) + +var newAIProgramFunc = func(model tea.Model) *tea.Program { + return tea.NewProgram(model, tea.WithAltScreen()) +} + +func NewAICommand() *cobra.Command { + var modelStr string + var baseURLStr string + var apiKeyStr string + var unsafeAllowWrite bool + + cmd := &cobra.Command{ + Use: "ai [PROMPT]", + Short: "Interactive AI assistant mode (TUI)", + RunE: func(cmd *cobra.Command, args []string) error { + prompt := "" + if len(args) > 0 { + prompt = strings.Join(args, " ") + } + + opts := config.Options{ + ConfigPath: GlobalConfig.ConfigStr, + CLIProfile: GlobalConfig.ProfileStr, + CLIProfileSet: cmd.Flags().Changed("profile"), + CLIAIModel: modelStr, + CLIAIModelSet: cmd.Flags().Changed("model"), + CLIAIBaseURL: baseURLStr, + CLIAIBaseURLSet: cmd.Flags().Changed("base-url"), + CLIAIAPIKey: apiKeyStr, + CLIAIAPIKeySet: cmd.Flags().Changed("api-key"), + } + + resolved, xe := config.Resolve(opts) + if xe != nil { + return xe + } + if resolved.ProfileName == "" || resolved.Profile.DB == "" { + return fmt.Errorf("config error [XSQL_CFG_INVALID]: no profile specified and no 'default' profile found in config") + } + + // Resolve API key if keyring reference or plaintext + apiKey := resolved.AI.APIKey + if secret.IsKeyringRef(apiKey) { + resolvedKey, xe := secret.Resolve(apiKey, secret.Options{AllowPlaintext: true}) + if xe == nil { + apiKey = resolvedKey + } + } + resolved.AI.APIKey = apiKey + + aiClient := ai.NewClient(resolved.AI, nil) + aiService := ai.NewService(resolved.AI, aiClient) + + m := tui.NewModel(opts, resolved, aiService, prompt, unsafeAllowWrite) + p := newAIProgramFunc(m) + if _, err := p.Run(); err != nil { + return fmt.Errorf("error running TUI: %w", err) + } + return nil + }, + } + + cmd.Flags().StringVar(&modelStr, "model", "", "AI model name (default: gpt-4o)") + cmd.Flags().StringVar(&baseURLStr, "base-url", "", "AI service base URL") + cmd.Flags().StringVar(&apiKeyStr, "api-key", "", "AI service API key") + cmd.Flags().BoolVar(&unsafeAllowWrite, "unsafe-allow-write", false, "Allow write operations (bypasses read-only protection)") + + return cmd +} diff --git a/cmd/xsql/ai_test.go b/cmd/xsql/ai_test.go new file mode 100644 index 0000000..386d7c9 --- /dev/null +++ b/cmd/xsql/ai_test.go @@ -0,0 +1,45 @@ +package main + +import ( + "os" + "path/filepath" + "strings" + "testing" + + tea "github.com/charmbracelet/bubbletea" +) + +func TestCmdXSQL_AICommand(t *testing.T) { + oldNewAIProgram := newAIProgramFunc + defer func() { newAIProgramFunc = oldNewAIProgram }() + + newAIProgramFunc = func(model tea.Model) *tea.Program { + p := tea.NewProgram(model, tea.WithInput(strings.NewReader("")), tea.WithOutput(os.Stderr), tea.WithoutRenderer()) + go p.Quit() + return p + } + + tmpDir := t.TempDir() + cfgPath := filepath.Join(tmpDir, "xsql.yaml") + cfgContent := ` +profiles: + dev: + db: mysql + host: 127.0.0.1 + port: 3306 + user: root + database: testdb +` + if err := os.WriteFile(cfgPath, []byte(cfgContent), 0600); err != nil { + t.Fatalf("failed to write temp config: %v", err) + } + + root := NewRootCommand() + root.AddCommand(NewAICommand()) + root.SetArgs([]string{"ai", "--config", cfgPath, "--profile", "dev", "Show top 10 users"}) + + err := root.Execute() + if err != nil { + t.Fatalf("expected xsql ai command execution to succeed, got %v", err) + } +} diff --git a/cmd/xsql/main.go b/cmd/xsql/main.go index 218da23..e50ad20 100644 --- a/cmd/xsql/main.go +++ b/cmd/xsql/main.go @@ -34,6 +34,7 @@ func run() int { root.AddCommand(NewServeCommand(&w)) root.AddCommand(NewWebCommand(&w)) root.AddCommand(NewStatsCommand(&w)) + root.AddCommand(NewAICommand()) // Execute and handle errors if err := root.Execute(); err != nil { diff --git a/cmd/xsql/root.go b/cmd/xsql/root.go index b229b27..20abf94 100644 --- a/cmd/xsql/root.go +++ b/cmd/xsql/root.go @@ -6,6 +6,8 @@ import ( "github.com/spf13/cobra" "github.com/zx06/xsql/internal/config" + _ "github.com/zx06/xsql/internal/db/mysql" + _ "github.com/zx06/xsql/internal/db/pg" "github.com/zx06/xsql/internal/errors" "github.com/zx06/xsql/internal/stats" ) diff --git a/docs/ai.md b/docs/ai.md index 0a31991..34082a2 100644 --- a/docs/ai.md +++ b/docs/ai.md @@ -87,7 +87,7 @@ MCP Server 提供以下 tools： "mcpServers": { "xsql": { "command": "xsql", - "args": ["mcp", "server", "--config", "/path/to/config.yaml"] + "args": ["mcp", "server", "--command", "mcp", "server", "--config", "/path/to/config.yaml"] } } } @@ -105,3 +105,45 @@ xsql web ``` Web UI 复用 xsql 的 profile、SSH、只读策略和结构化错误契约，但其 HTTP API 面向浏览器，不等同于 MCP 协议。 + +## AI TUI 交互模式 (xsql-ai) +xsql-ai 为独立的 CLI 可执行程序，提供交互式 AI 终端模式。用户只需在终端以自然语言发问，AI 结合当前数据库 Schema 结构自动构建对应的 SQL 查询，并在 TUI 中提供交互预览与安全执行： + +```bash +# 启动交互式 TUI +xsql-ai --profile dev +``` + +### LLM 集成与 Tool Call 机制 +`xsql` 使用 OpenAI 官方 SDK (`github.com/openai/openai-go`) 与大模型交互，基于标准的 **ReAct Agent Loop 循环推理**，支持 3 大核心 Tools 调度： +1. **`execute_sql(sql: string, explanation: string)`**: 数据库 SQL 查询工具（执行成功后宿主自动渲染内嵌交互表格）。 +2. **`execute_javascript(js_code: string, explanation: string)`**: 基于 `goja` 沙箱的本地 JS 数据聚合计算工具（必须遵循 ES5 语法）。 +3. **`export_data(dataset_id: string, format: string, filepath: string, explanation: string)`**: 会话数据集文件导出工具（触发人机交互二次确认）。 + +#### ReAct Agent Loop 准则 +- **循环驱动**：Agent 会在单次交互中循环执行 Tools，直到不再产生 Tool Call。 +- **最终回答不变性**：交互轮次的最终输出必定是 AI 总结出的自然语言 / Markdown 格式分析报告。 +- **工具折叠与容器内嵌**：所有的中间 Tool Call 默认以单行 Pill 收起折叠（内嵌表格与指标数据），界面保持极简清爽。 + +#### 零数据泄露与 Session 数据集召回 (Session DataStore) +- 每次查询成功的结果在本地分配标号（`res1`, `res2`, ...）。 +- 大模型上下文中仅包含数据集的轻量 Catalog 目录结构（字段名与行数），不传输海量真实数据。 +- AI 可通过 `execute_javascript` 生成纯 Go 沙箱 (`goja`) 执行的代码，在本地对 `res1`, `res2` 等数据集做跨表 Join、占比统计与数据清洗，并通过 `export_data` 安全导出为 CSV/JSON/Markdown。 + +### 快捷键操作 + +#### SQL & 导出确认状态 (Approval Mode) +- `Enter`: 确认并安全执行当前生成预览的 SQL 或同意文件导出 +- `e`: 切换到 SQL 文本手工编辑/微调模式 +- `Esc`: 取消当前 SQL 生成建议或拒绝文件导出 + +#### 通用与表格/工具操作 (General & Tool Operations) +- `Enter`: 提交自然语言需求给 AI +- `Ctrl+O`: 折叠/展开当前选中的 Tool Call 详情（内嵌表格与指标） +- `Ctrl+P` / `Ctrl+N`: 在会话历史中的多个 Tool Call 组件之间向前/向后切换焦点 +- `Ctrl+E`: 切换表格单行展开视图（Expand Vertical View） +- `Tab`: 在历史多个查询结果表格之间无缝切换焦点 (`[FOCUSED]`) +- `←` / `→`: 横向平滑滚动查看当前焦点表格的隐藏列 +- `PgUp` / `PgDn`: 向上/向下翻页查看当前焦点表格的数据 +- `Shift+Tab`: 一键切换 **自动执行 (AUTO-EXECUTE)** 与 **手动批准 (MANUAL-APPROVE)** 模式 +- `Esc` / `Ctrl+C`: 退出 AI 模式 diff --git a/docs/architecture.md b/docs/architecture.md index 5e640e1..b590660 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -6,9 +6,12 @@ ## 建议目录结构 ``` -/cmd/xsql # CLI 入口 +/cmd/xsql # CLI 主程序入口 +/cmd/xsql-ai # AI TUI 独立程序入口 /internal/app # 应用编排（解析参数->执行->返回结构化结果） /internal/config # 配置加载/合并/校验 + profiles +/internal/ai # AI LLM 客户端与 Context Prompt 组装 +/internal/tui # Bubbletea TUI 交互式终端UI实现 /internal/secret # keyring/加密/明文兼容 /internal/db # driver registry + 执行引擎 /internal/db/mysql # MySQL 驱动实现 diff --git a/docs/cli-spec.md b/docs/cli-spec.md index 5a96834..1bde715 100644 --- a/docs/cli-spec.md +++ b/docs/cli-spec.md @@ -546,6 +546,35 @@ xsql mcp server --transport streamable_http --http-addr 127.0.0.1:8787 --http-au - 写操作需要显式设置 `unsafe_allow_write: true` - Streamable HTTP 传输要求鉴权，请在请求中提供 `Authorization: Bearer ` 头 +### `xsql-ai` 独立程序 + +`xsql-ai` 为独立的 CLI 可执行程序，提供类似 Chatbot 的交互终端（TUI）。通过自然语言与 AI 对话，由 AI 基于当前数据库的 Schema 结构自动构建 SQL 查询，并在终端进行可视化预览与安全执行。 + +```bash +# 启动交互式 TUI 模式 +xsql-ai --profile dev + +# 指定 AI 模型和服务地址 +xsql-ai --profile dev --model deepseek-coder --base-url https://api.deepseek.com/v1 + +# 启动并直接传入初始 Prompt 自动分析执行 +xsql-ai --profile dev "查一下近7天注册的用户数量" +``` + +**Flags:** +| Flag | 默认值 | 说明 | +|------|--------|------| +| `--profile` | - | Profile 名称（必需） | +| `--model` | `gpt-4o` | AI 模型名称（配置项：`ai.model`，ENV：`XSQL_AI_MODEL`） | +| `--base-url` | `https://api.openai.com/v1` | AI 服务端 Base URL（配置项：`ai.base_url`，ENV：`XSQL_AI_BASE_URL`） | +| `--api-key` | - | AI 服务 API Key（配置项：`ai.api_key`，ENV：`XSQL_AI_API_KEY`） | +| `--prompt` | - | 单次自然语言提问 Prompt（也可通过命令行位置参数传入） | +| `--unsafe-allow-write` | false | 允许写操作（绕过只读保护） | + +**安全与只读说明：** +- 默认启用只读 protection。AI 生成的 SQL 语句在运行（按 `Ctrl+E`）时将统一提交由底层只读策略审计及事务级只读操作校验。 +- API Key 在配置中支持 `keyring:ai_key` 密码存储。 + ## 参数来源优先级 - CLI > ENV > Config diff --git a/docs/rfcs/0008-ai-tui.md b/docs/rfcs/0008-ai-tui.md new file mode 100644 index 0000000..65b2c0d --- /dev/null +++ b/docs/rfcs/0008-ai-tui.md @@ -0,0 +1,54 @@ +# RFC 0008: AI TUI Interactive Mode and Standalone `xsql-ai` Binary + +Status: Accepted + +## 摘要 +本 RFC 提出为 `xsql` 增加交互式 AI 模式（TUI）及独立终端可执行程序 `cmd/xsql-ai`。用户可以通过自然语言在终端提问，AI 自动结合当前数据库的 Schema 结构生成相应的 SQL 语句，并在 TUI 界面中提供可视化预览、编辑与一键安全执行。 + +## 背景 / 动机 +- **当前痛点**：用户在面对复杂的数据库表结构时，手写 SQL 查询门槛较高，传统 CLI 查询缺少交互式的自然语言转 SQL 辅助。 +- **目标**： + 1. 提供独立终端程序 `xsql-ai` 及主程序子命令 `xsql ai`。 + 2. 支持标准 OpenAI 兼容接口配置（兼容 OpenAI, DeepSeek, Ollama 等）。 + 3. 自动抽取 Profile 对应的数据库 Schema，作为上下文送给 LLM 生成对应 DB 语法的 SQL。 + 4. 采用 Charm (Bubbletea / Lipgloss) 编写现代、美观、响应式的 TUI 界面，支持 SQL 预览、编辑与快捷执行。 + 5. 严格保留 `xsql` 既有的只读策略保护（双重只读拦截：SQL 静态检测 + 事务级 READ ONLY）。 +- **非目标**： + - 本 RFC 不试图在 CLI 内实现大模型训练或复杂的全表数据检索，仅聚焦于辅助 SQL 编写与交互式查询执行。 + +## 方案（Proposed） + +### 用户视角（CLI/配置/输出） +1. **独立程序与 CLI 命令**： + - 独立可执行程序：`xsql-ai -p [--prompt "query"]` + - 主程序子命令：`xsql ai -p ` +2. **配置文件扩展**： + 在 `xsql.yaml` 中新增 `ai` 配置块： + ```yaml + ai: + provider: openai + base_url: https://api.openai.com/v1 + api_key: keyring:ai_key + model: gpt-4o + max_tokens: 2048 + ``` +3. **优先级与凭据**： + - 合并优先级：`CLI flags > ENV (XSQL_AI_API_KEY, XSQL_AI_BASE_URL, XSQL_AI_MODEL) > Config`。 + - API Key 支持 `keyring:` 引用。 + +### 技术设计（Architecture） +- **涉及模块**： + - `cmd/xsql-ai/main.go`：独立 CLI 程序入口。 + - `cmd/xsql/ai.go`：`xsql ai` 适配入口。 + - `internal/config`：配置类型扩展与解析。 + - `internal/ai`：OpenAI HTTP 客户端、Prompt 组装服务。 + - `internal/tui`：Bubbletea UI 架构（Header、Viewport、SQL Preview Card、Textarea）。 + - `internal/app`：复用 `DumpSchema` 与 `Query`。 + +### 安全与隐私（Security/Privacy） +- 默认严格只读策略，除非显式设置 `--unsafe-allow-write`，否则不允许写 SQL 执行。 +- 不会把明文 API Key 输出在日志或错细节中。 + +### 测试计划（Test Plan） +- **单元测试**：配置合并解析、Prompt 组装、OpenAI 响应解析、TUI 状态机逻辑测试。 +- **E2E 测试**：在 `tests/e2e/ai_test.go` 中启动 Mock HTTP API Server 并通过 IO Pipe 输入按键模拟全闭环交互。 diff --git a/docs/rfcs/0009-openai-sdk-tool-call.md b/docs/rfcs/0009-openai-sdk-tool-call.md new file mode 100644 index 0000000..c699d91 --- /dev/null +++ b/docs/rfcs/0009-openai-sdk-tool-call.md @@ -0,0 +1,35 @@ +# RFC 0009: Migrate OpenAI Integration to Official SDK and Tool Call + +Status: Proposed + +## 摘要 +本 RFC 提出将 `xsql` 项目中的 AI 客户端重构为使用 OpenAI 官方 Go SDK (`github.com/openai/openai-go`)，并将 SQL 生成与解析机制从手动字符串/JSON 解析改为标准 OpenAI Tool Call (`execute_sql`) 模式。 + +## 背景 / 动机 +- **当前问题**： + 1. `internal/ai/client.go` 自行实现了 HTTP 请求封装，维护成本高且不易支持高级特性。 + 2. `internal/ai/service.go` 通过 System Prompt 强求模型输出 JSON 字符串，并通过正则表达式/`json.Unmarshal` 提取 `sql` 和 `explanation`，解析脆弱且容易由于 markdown 格式干扰出错。 +- **目标**： + 1. 引入 `github.com/openai/openai-go` 官方 SDK。 + 2. 定义 `execute_sql(sql string, explanation string)` 函数工具，由模型通过 Tool Call 返回结构化参数。 + 3. 保留对无 Tool Call 场景（文本回复）的兼容容错。 + +## 方案（Proposed） + +### 技术设计 +1. **官方 SDK 集成**： + - 依赖：`github.com/openai/openai-go` + - 初始化：使用 `openai.NewClient(option.WithAPIKey(...), option.WithBaseURL(...), option.WithHTTPClient(...))`。 +2. **Tool Calling 定义**： + - 工具名称：`execute_sql` + - 工具描述：Execute or present generated SQL query based on database schema and user intent. + - 参数 Schema： + - `sql`: (string) 针对特定 DB 语法的 SQL 查询语句。 + - `explanation`: (string) 对查询意图或无法生成 SQL 的说明解释。 +3. **响应处理逻辑**： + - 若模型响应包含 `execute_sql` 的 Tool Call，则解析 JSON 参数提取 `sql` 与 `explanation`。 + - 若模型仅返回文本（无 Tool Call），则设置 `sql=""`，并将文本存入 `explanation`。 + +### 测试计划 +- 单元测试：`internal/ai/service_test.go` 升级 Mock API 响应为 Tool Call 消息体格式。 +- E2E 测试：`tests/e2e/ai_test.go` 升级 Mock API 响应。 diff --git a/docs/rfcs/0010-tui-visual-redesign.md b/docs/rfcs/0010-tui-visual-redesign.md new file mode 100644 index 0000000..220a82f --- /dev/null +++ b/docs/rfcs/0010-tui-visual-redesign.md @@ -0,0 +1,30 @@ +# RFC 0010: TUI Visual and Interactive Experience Redesign + +Status: Proposed + +## 摘要 +本 RFC 提出对 `xsql-ai` 终端 TUI 进行视觉与交互体验的重构。参考 Catppuccin / Tokyo Night 现代调色盘以及 Charm (Bubbletea / Lipgloss) 社区开源工具（如 `mods`、`glow`、`gh-dash`）的设计精髓，全面优化 Pill 胶囊标签、✦ Prompt 输入框、全主题自适应调色盘与按键指引。 + +## 背景 / 动机 +- **当前问题**： + - Header 与提示文本在白色/浅色终端主题下对比度不理想。 + - 缺少 SQL 执行的耗时与元信息状态展示。 + - 底部快捷键文本较为平淡，缺乏现代终端工具的按键 Badge 指示器。 + +## 方案（Proposed） + +### 视觉与交互规范 +1. **Pill 胶囊 Header 导航**： + - 使用 rounded 内边距与不同主题色背景塑造 `xsql AI`、Profile、DB、只读/读写模式与自动/手动执行模式。 +2. **运行指标 (Metrics Bar)**： + - 执行 SQL 后输出包含执行耗时（如 `⏱️ 14ms`）、行数（如 `📊 10 rows`）与 LLM 模型（如 `🤖 gpt-4o`）。 +3. **✦ Prompt 输入框与 Focus 指示器**： + - 输入框增加 `✦ Ask AI:` 品牌提示前缀，并根据获得焦点状态显示鲜明边框。 +4. **按键 Badge 底部栏**： + - 底部提示升级为形如 `[Enter] 发送` `[Tab] 聚焦表格` `[e] 编辑SQL` `[Shift+Tab] 模式切换` `[Esc] 退出` 的精致键盘 Badge。 +5. **Catppuccin / Tokyo Night 调色盘**： + - 全量采用 `lipgloss.AdaptiveColor` 确保在任何深色/浅色背景终端中均具有符合 WCAG 的高对比度。 + +### 测试计划 +- 单元测试：`internal/tui/components_test.go` 与 `model_test.go`。 +- E2E 测试：`tests/e2e/ai_test.go`。 diff --git a/docs/rfcs/0011-goja-js-data-analysis.md b/docs/rfcs/0011-goja-js-data-analysis.md new file mode 100644 index 0000000..4c75e37 --- /dev/null +++ b/docs/rfcs/0011-goja-js-data-analysis.md @@ -0,0 +1,33 @@ +# RFC 0011: Integration of goja JS Engine, Session DataStore, and ReAct Tool Agent Loop + +Status: Proposed + +## 摘要 +本 RFC 提出在 `xsql` / `xsql-ai` 中集成纯 Go 实现的 `goja` JavaScript 虚拟机（100% Zero CGO），构建 **Session DataStore（会话数据集存储与召回）** 机制，并实现无硬编码的 **ReAct Tool Agent Loop 循环推理**。 + +## 背景 / 动机 +- 当前 `xsql-ai` 仅支持 SQL 交互与表格展示，缺少数据二次聚合计算、跨查询结果 Join/比对以及结构化导出文件（CSV/JSON/Markdown）的能力。 +- 将海量原始数据全量透传给大模型（LLM）会导致上下文爆炸（Context Overflow）与高昂 Token 成本，且存在数据隐私泄露红线。 + +## 架构与核心设计 + +### 1. 零 CGO JS 引擎 (`internal/js`) +- 使用 `github.com/dop251/goja` 在纯 Go 内存沙箱中执行 AI 动态生成的 JS 数据分析代码。 +- 规定 JS 代码必须遵循 ES5 (ECMAScript 5.1) 标准语法。 +- 支持 Context 超时打断（默认 1 分钟，可配置 `js_timeout`），并自动捕获 `console.log` 输出。 +- 包含 AI 自动重试修正机制（上限 3 次）。 + +### 2. Session 数据集存储与召回 (`internal/session`) +- 本地维护 `SessionDataStore`，为每次 SQL 执行成功的 QueryResult 分配唯一 ID（`res1`, `res2`, ...）。 +- 向大模型上下文仅提供轻量 **Dataset Catalog** 元数据目录，LLM 可以在后续多轮对话中指定 `res1`, `res2` 召回历史数据并在 JS 中做跨数据集 Join 或计算。 + +### 3. 外层文件导出 (`internal/export`) +- JS 仅负责数据计算与转换；由外层 Go 宿主层统一执行安全的磁盘文件写入（CSV / JSON / Markdown），并强制人机交互二次确认（Human-in-the-loop）。 + +### 4. ReAct Tool Agent Loop 架构 (`internal/ai` & `internal/tui`) +AI 具备 3 大解耦工具： +1. `execute_sql`: 数据库 SQL 查询（执行完由宿主层自动渲染内嵌交互表格） +2. `execute_javascript`: ES5 沙箱数据二次清洗与聚合 +3. `export_data`: 文件导出（含人机交互确认卡片） + +所有 Tool Calls 默认在 TUI 容器中折叠内嵌呈现（`Ctrl+O` 展开/折叠，`Ctrl+P`/`Ctrl+N` 切换焦点），且交互末尾必定以 LLM 自然语言 Markdown 分析报告总结收尾。 diff --git a/go.mod b/go.mod index b58e683..8b9b039 100644 --- a/go.mod +++ b/go.mod @@ -3,10 +3,15 @@ module github.com/zx06/xsql go 1.25.0 require ( + github.com/charmbracelet/bubbles v0.20.0 + github.com/charmbracelet/bubbletea v1.3.4 + github.com/charmbracelet/lipgloss v1.1.1-0.20250404203927-76690c660834 github.com/go-sql-driver/mysql v1.10.0 github.com/google/jsonschema-go v0.4.3 github.com/jackc/pgx/v5 v5.9.2 + github.com/mattn/go-runewidth v0.0.17 github.com/modelcontextprotocol/go-sdk v1.6.0 + github.com/openai/openai-go v1.12.0 github.com/spf13/cobra v1.10.2 github.com/zalando/go-keyring v0.2.8 golang.org/x/crypto v0.52.0 @@ -17,18 +22,52 @@ require ( require ( filippo.io/edwards25519 v1.2.0 // indirect + github.com/alecthomas/chroma/v2 v2.20.0 // indirect + github.com/atotto/clipboard v0.1.4 // indirect + github.com/aymanbagabas/go-osc52/v2 v2.0.1 // indirect + github.com/aymerick/douceur v0.2.0 // indirect + github.com/charmbracelet/colorprofile v0.2.3-0.20250311203215-f60798e515dc // indirect + github.com/charmbracelet/glamour v1.0.0 // indirect + github.com/charmbracelet/x/ansi v0.10.2 // indirect + github.com/charmbracelet/x/cellbuf v0.0.13 // indirect + github.com/charmbracelet/x/exp/slice v0.0.0-20250327172914-2fdc97757edf // indirect + github.com/charmbracelet/x/term v0.2.1 // indirect github.com/danieljoos/wincred v1.2.3 // indirect + github.com/dlclark/regexp2 v1.11.5 // indirect + github.com/dlclark/regexp2/v2 v2.5.2 // indirect + github.com/dop251/goja v0.0.0-20260723142020-b4aef50fa347 // indirect + github.com/erikgeiser/coninput v0.0.0-20211004153227-1c3628e74d0f // indirect + github.com/go-sourcemap/sourcemap v2.1.3+incompatible // indirect github.com/godbus/dbus/v5 v5.2.2 // indirect + github.com/google/pprof v0.0.0-20230207041349-798e818bf904 // indirect + github.com/gorilla/css v1.0.1 // indirect github.com/inconshreveable/mousetrap v1.1.0 // indirect github.com/jackc/pgpassfile v1.0.0 // indirect github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 // indirect github.com/jackc/puddle/v2 v2.2.2 // indirect github.com/kr/text v0.2.0 // indirect + github.com/lucasb-eyer/go-colorful v1.3.0 // indirect + github.com/mattn/go-isatty v0.0.20 // indirect + github.com/mattn/go-localereader v0.0.1 // indirect + github.com/microcosm-cc/bluemonday v1.0.27 // indirect + github.com/muesli/ansi v0.0.0-20230316100256-276c6243b2f6 // indirect + github.com/muesli/cancelreader v0.2.2 // indirect + github.com/muesli/reflow v0.3.0 // indirect + github.com/muesli/termenv v0.16.0 // indirect + github.com/rivo/uniseg v0.4.7 // indirect github.com/rogpeppe/go-internal v1.14.1 // indirect github.com/segmentio/asm v1.2.1 // indirect github.com/segmentio/encoding v0.5.4 // indirect github.com/spf13/pflag v1.0.10 // indirect + github.com/tidwall/gjson v1.14.4 // indirect + github.com/tidwall/match v1.1.1 // indirect + github.com/tidwall/pretty v1.2.1 // indirect + github.com/tidwall/sjson v1.2.5 // indirect + github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e // indirect github.com/yosida95/uritemplate/v3 v3.0.2 // indirect + github.com/yuin/goldmark v1.7.13 // indirect + github.com/yuin/goldmark-emoji v1.0.6 // indirect + golang.org/x/net v0.54.0 // indirect golang.org/x/oauth2 v0.36.0 // indirect golang.org/x/sys v0.45.0 // indirect golang.org/x/text v0.37.0 // indirect diff --git a/go.sum b/go.sum index 153b2a8..53dac4c 100644 --- a/go.sum +++ b/go.sum @@ -1,5 +1,41 @@ filippo.io/edwards25519 v1.2.0 h1:crnVqOiS4jqYleHd9vaKZ+HKtHfllngJIiOpNpoJsjo= filippo.io/edwards25519 v1.2.0/go.mod h1:xzAOLCNug/yB62zG1bQ8uziwrIqIuxhctzJT18Q77mc= +github.com/MakeNowJust/heredoc v1.0.0 h1:cXCdzVdstXyiTqTvfqk9SDHpKNjxuom+DOlyEeQ4pzQ= +github.com/MakeNowJust/heredoc v1.0.0/go.mod h1:mG5amYoWBHf8vpLOuehzbGGw0EHxpZZ6lCpQ4fNJ8LE= +github.com/alecthomas/chroma/v2 v2.20.0 h1:sfIHpxPyR07/Oylvmcai3X/exDlE8+FA820NTz+9sGw= +github.com/alecthomas/chroma/v2 v2.20.0/go.mod h1:e7tViK0xh/Nf4BYHl00ycY6rV7b8iXBksI9E359yNmA= +github.com/atotto/clipboard v0.1.4 h1:EH0zSVneZPSuFR11BlR9YppQTVDbh5+16AmcJi4g1z4= +github.com/atotto/clipboard v0.1.4/go.mod h1:ZY9tmq7sm5xIbd9bOK4onWV4S6X0u6GY7Vn0Yu86PYI= +github.com/aymanbagabas/go-osc52/v2 v2.0.1 h1:HwpRHbFMcZLEVr42D4p7XBqjyuxQH5SMiErDT4WkJ2k= +github.com/aymanbagabas/go-osc52/v2 v2.0.1/go.mod h1:uYgXzlJ7ZpABp8OJ+exZzJJhRNQ2ASbcXHWsFqH8hp8= +github.com/aymanbagabas/go-udiff v0.2.0 h1:TK0fH4MteXUDspT88n8CKzvK0X9O2xu9yQjWpi6yML8= +github.com/aymanbagabas/go-udiff v0.2.0/go.mod h1:RE4Ex0qsGkTAJoQdQQCA0uG+nAzJO/pI/QwceO5fgrA= +github.com/aymerick/douceur v0.2.0 h1:Mv+mAeH1Q+n9Fr+oyamOlAkUNPWPlA8PPGR0QAaYuPk= +github.com/aymerick/douceur v0.2.0/go.mod h1:wlT5vV2O3h55X9m7iVYN0TBM0NH/MmbLnd30/FjWUq4= +github.com/charmbracelet/bubbles v0.20.0 h1:jSZu6qD8cRQ6k9OMfR1WlM+ruM8fkPWkHvQWD9LIutE= +github.com/charmbracelet/bubbles v0.20.0/go.mod h1:39slydyswPy+uVOHZ5x/GjwVAFkCsV8IIVy+4MhzwwU= +github.com/charmbracelet/bubbletea v1.3.4 h1:kCg7B+jSCFPLYRA52SDZjr51kG/fMUEoPoZrkaDHyoI= +github.com/charmbracelet/bubbletea v1.3.4/go.mod h1:dtcUCyCGEX3g9tosuYiut3MXgY/Jsv9nKVdibKKRRXo= +github.com/charmbracelet/colorprofile v0.2.3-0.20250311203215-f60798e515dc h1:4pZI35227imm7yK2bGPcfpFEmuY1gc2YSTShr4iJBfs= +github.com/charmbracelet/colorprofile v0.2.3-0.20250311203215-f60798e515dc/go.mod h1:X4/0JoqgTIPSFcRA/P6INZzIuyqdFY5rm8tb41s9okk= +github.com/charmbracelet/glamour v1.0.0 h1:AWMLOVFHTsysl4WV8T8QgkQ0s/ZNZo7CiE4WKhk8l08= +github.com/charmbracelet/glamour v1.0.0/go.mod h1:DSdohgOBkMr2ZQNhw4LZxSGpx3SvpeujNoXrQyH2hxo= +github.com/charmbracelet/lipgloss v1.0.0 h1:O7VkGDvqEdGi93X+DeqsQ7PKHDgtQfF8j8/O2qFMQNg= +github.com/charmbracelet/lipgloss v1.0.0/go.mod h1:U5fy9Z+C38obMs+T+tJqst9VGzlOYGj4ri9reL3qUlo= +github.com/charmbracelet/lipgloss v1.1.1-0.20250404203927-76690c660834 h1:ZR7e0ro+SZZiIZD7msJyA+NjkCNNavuiPBLgerbOziE= +github.com/charmbracelet/lipgloss v1.1.1-0.20250404203927-76690c660834/go.mod h1:aKC/t2arECF6rNOnaKaVU6y4t4ZeHQzqfxedE/VkVhA= +github.com/charmbracelet/x/ansi v0.8.0 h1:9GTq3xq9caJW8ZrBTe0LIe2fvfLR/bYXKTx2llXn7xE= +github.com/charmbracelet/x/ansi v0.8.0/go.mod h1:wdYl/ONOLHLIVmQaxbIYEC/cRKOQyjTkowiI4blgS9Q= +github.com/charmbracelet/x/ansi v0.10.2 h1:ith2ArZS0CJG30cIUfID1LXN7ZFXRCww6RUvAPA+Pzw= +github.com/charmbracelet/x/ansi v0.10.2/go.mod h1:HbLdJjQH4UH4AqA2HpRWuWNluRE6zxJH/yteYEYCFa8= +github.com/charmbracelet/x/cellbuf v0.0.13 h1:/KBBKHuVRbq1lYx5BzEHBAFBP8VcQzJejZ/IA3iR28k= +github.com/charmbracelet/x/cellbuf v0.0.13/go.mod h1:xe0nKWGd3eJgtqZRaN9RjMtK7xUYchjzPr7q6kcvCCs= +github.com/charmbracelet/x/exp/golden v0.0.0-20240815200342-61de596daa2b h1:MnAMdlwSltxJyULnrYbkZpp4k58Co7Tah3ciKhSNo0Q= +github.com/charmbracelet/x/exp/golden v0.0.0-20240815200342-61de596daa2b/go.mod h1:wDlXFlCrmJ8J+swcL/MnGUuYnqgQdW9rhSD61oNMb6U= +github.com/charmbracelet/x/exp/slice v0.0.0-20250327172914-2fdc97757edf h1:rLG0Yb6MQSDKdB52aGX55JT1oi0P0Kuaj7wi1bLUpnI= +github.com/charmbracelet/x/exp/slice v0.0.0-20250327172914-2fdc97757edf/go.mod h1:B3UgsnsBZS/eX42BlaNiJkD1pPOUa+oF1IYC6Yd2CEU= +github.com/charmbracelet/x/term v0.2.1 h1:AQeHeLZ1OqSXhrAWpYUtZyX1T3zVxfpZuEQMIQaGIAQ= +github.com/charmbracelet/x/term v0.2.1/go.mod h1:oQ4enTYFV7QN4m0i9mzHrViD7TQKvNEEkHUMCmsxdUg= github.com/cpuguy83/go-md2man/v2 v2.0.6/go.mod h1:oOW0eioCTA6cOiMLiUPZOpcVxMig6NIQQ7OS05n1F4g= github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E= github.com/danieljoos/wincred v1.2.3 h1:v7dZC2x32Ut3nEfRH+vhoZGvN72+dQ/snVXo/vMFLdQ= @@ -7,6 +43,16 @@ github.com/danieljoos/wincred v1.2.3/go.mod h1:6qqX0WNrS4RzPZ1tnroDzq9kY3fu1KwE7 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/dlclark/regexp2 v1.11.5 h1:Q/sSnsKerHeCkc/jSTNq1oCm7KiVgUMZRDUoRu0JQZQ= +github.com/dlclark/regexp2 v1.11.5/go.mod h1:DHkYz0B9wPfa6wondMfaivmHpzrQ3v9q8cnmRbL6yW8= +github.com/dlclark/regexp2/v2 v2.5.2 h1:HAsucWRhsqcDzl6Ua9aR8JwYOTzrZyPrF0/FNxJVAI0= +github.com/dlclark/regexp2/v2 v2.5.2/go.mod h1:avUrQvPaLz2DrFNHJF0taWAFFX2C1GMSSoeiqFjcBmU= +github.com/dop251/goja v0.0.0-20260723142020-b4aef50fa347 h1:RZr+96+PKQjn444QL1K9MtncwJ/PwfE+3TJLCYJL8es= +github.com/dop251/goja v0.0.0-20260723142020-b4aef50fa347/go.mod h1:LiIEzozrcvNXorsG/3+ypGqdTUAqZryhzSsqi0oU/Qg= +github.com/erikgeiser/coninput v0.0.0-20211004153227-1c3628e74d0f h1:Y/CXytFA4m6baUTXGLOoWe4PQhGxaX0KpnayAqC48p4= +github.com/erikgeiser/coninput v0.0.0-20211004153227-1c3628e74d0f/go.mod h1:vw97MGsxSvLiUE2X8qFplwetxpGLQrlU1Q9AUEIzCaM= +github.com/go-sourcemap/sourcemap v2.1.3+incompatible h1:W1iEw64niKVGogNgBN3ePyLFfuisuzeidWPMPWmECqU= +github.com/go-sourcemap/sourcemap v2.1.3+incompatible/go.mod h1:F8jJfvm2KbVjc5NqelyYJmf/v5J0dwNLS2mL4sNA1Jg= github.com/go-sql-driver/mysql v1.10.0 h1:Q+1LV8DkHJvSYAdR83XzuhDaTykuDx0l6fkXxoWCWfw= github.com/go-sql-driver/mysql v1.10.0/go.mod h1:M+cqaI7+xxXGG9swrdeUIoPG3Y3KCkF0pZej+SK+nWk= github.com/godbus/dbus/v5 v5.2.2 h1:TUR3TgtSVDmjiXOgAAyaZbYmIeP3DPkld3jgKGV8mXQ= @@ -17,6 +63,10 @@ github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8= github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU= github.com/google/jsonschema-go v0.4.3 h1:/DBOLZTfDow7pe2GmaJNhltueGTtDKICi8V8p+DQPd0= github.com/google/jsonschema-go v0.4.3/go.mod h1:r5quNTdLOYEz95Ru18zA0ydNbBuYoo9tgaYcxEYhJVE= +github.com/google/pprof v0.0.0-20230207041349-798e818bf904 h1:4/hN5RUoecvl+RmJRE2YxKWtnnQls6rQjjW5oV7qg2U= +github.com/google/pprof v0.0.0-20230207041349-798e818bf904/go.mod h1:uglQLonpP8qtYCYyzA+8c/9qtqgA3qsXGYqCPKARAFg= +github.com/gorilla/css v1.0.1 h1:ntNaBIghp6JmvWnxbZKANoLyuXTPZ4cAMlo6RyhlbO8= +github.com/gorilla/css v1.0.1/go.mod h1:BvnYkspnSzMmwRK+b8/xgNPLiIuNZr6vbZBTPQ2A3b0= github.com/inconshreveable/mousetrap v1.1.0 h1:wN+x4NVGpMsO7ErUn/mUI3vEoE6Jt13X2s0bqwp9tc8= github.com/inconshreveable/mousetrap v1.1.0/go.mod h1:vpF70FUmC8bwa3OWnCshd2FqLfsEA9PFc4w1p2J65bw= github.com/jackc/pgpassfile v1.0.0 h1:/6Hmqy13Ss2zCq62VdNG8tM1wchn8zjSGOBJ6icpsIM= @@ -31,10 +81,41 @@ github.com/kr/pretty v0.3.0 h1:WgNl7dwNpEZ6jJ9k1snq4pZsg7DOEN8hP9Xw0Tsjwk0= github.com/kr/pretty v0.3.0/go.mod h1:640gp4NfQd8pI5XOwp5fnNeVWj67G7CFk/SaSQn7NBk= github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY= github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE= +github.com/lucasb-eyer/go-colorful v1.2.0 h1:1nnpGOrhyZZuNyfu1QjKiUICQ74+3FNCN69Aj6K7nkY= +github.com/lucasb-eyer/go-colorful v1.2.0/go.mod h1:R4dSotOR9KMtayYi1e77YzuveK+i7ruzyGqttikkLy0= +github.com/lucasb-eyer/go-colorful v1.3.0 h1:2/yBRLdWBZKrf7gB40FoiKfAWYQ0lqNcbuQwVHXptag= +github.com/lucasb-eyer/go-colorful v1.3.0/go.mod h1:R4dSotOR9KMtayYi1e77YzuveK+i7ruzyGqttikkLy0= +github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY= +github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y= +github.com/mattn/go-localereader v0.0.1 h1:ygSAOl7ZXTx4RdPYinUpg6W99U8jWvWi9Ye2JC/oIi4= +github.com/mattn/go-localereader v0.0.1/go.mod h1:8fBrzywKY7BI3czFoHkuzRoWE9C+EiG4R1k4Cjx5p88= +github.com/mattn/go-runewidth v0.0.12/go.mod h1:RAqKPSqVFrSLVXbA8x7dzmKdmGzieGRCM46jaSJTDAk= +github.com/mattn/go-runewidth v0.0.16 h1:E5ScNMtiwvlvB5paMFdw9p4kSQzbXFikJ5SQO6TULQc= +github.com/mattn/go-runewidth v0.0.16/go.mod h1:Jdepj2loyihRzMpdS35Xk/zdY8IAYHsh153qUoGf23w= +github.com/mattn/go-runewidth v0.0.17 h1:78v8ZlW0bP43XfmAfPsdXcoNCelfMHsDmd/pkENfrjQ= +github.com/mattn/go-runewidth v0.0.17/go.mod h1:Jdepj2loyihRzMpdS35Xk/zdY8IAYHsh153qUoGf23w= +github.com/microcosm-cc/bluemonday v1.0.27 h1:MpEUotklkwCSLeH+Qdx1VJgNqLlpY2KXwXFM08ygZfk= +github.com/microcosm-cc/bluemonday v1.0.27/go.mod h1:jFi9vgW+H7c3V0lb6nR74Ib/DIB5OBs92Dimizgw2cA= github.com/modelcontextprotocol/go-sdk v1.6.0 h1:PPLS3kn7WtOEnR+Af4X5H96SG0qSab8R/ZQT/HkhPkY= github.com/modelcontextprotocol/go-sdk v1.6.0/go.mod h1:kzm3kzFL1/+AziGOE0nUs3gvPoNxMCvkxokMkuFapXQ= +github.com/muesli/ansi v0.0.0-20230316100256-276c6243b2f6 h1:ZK8zHtRHOkbHy6Mmr5D264iyp3TiX5OmNcI5cIARiQI= +github.com/muesli/ansi v0.0.0-20230316100256-276c6243b2f6/go.mod h1:CJlz5H+gyd6CUWT45Oy4q24RdLyn7Md9Vj2/ldJBSIo= +github.com/muesli/cancelreader v0.2.2 h1:3I4Kt4BQjOR54NavqnDogx/MIoWBFa0StPA8ELUXHmA= +github.com/muesli/cancelreader v0.2.2/go.mod h1:3XuTXfFS2VjM+HTLZY9Ak0l6eUKfijIfMUZ4EgX0QYo= +github.com/muesli/reflow v0.3.0 h1:IFsN6K9NfGtjeggFP+68I4chLZV2yIKsXJFNZ+eWh6s= +github.com/muesli/reflow v0.3.0/go.mod h1:pbwTDkVPibjO2kyvBQRBxTWEEGDGq0FlB1BIKtnHY/8= +github.com/muesli/termenv v0.15.2 h1:GohcuySI0QmI3wN8Ok9PtKGkgkFIk7y6Vpb5PvrY+Wo= +github.com/muesli/termenv v0.15.2/go.mod h1:Epx+iuz8sNs7mNKhxzH4fWXGNpZwUaJKRS1noLXviQ8= +github.com/muesli/termenv v0.16.0 h1:S5AlUN9dENB57rsbnkPyfdGuWIlkmzJjbFf0Tf5FWUc= +github.com/muesli/termenv v0.16.0/go.mod h1:ZRfOIKPFDYQoDFF4Olj7/QJbW60Ol/kL1pU3VfY/Cnk= +github.com/openai/openai-go v1.12.0 h1:NBQCnXzqOTv5wsgNC36PrFEiskGfO5wccfCWDo9S1U0= +github.com/openai/openai-go v1.12.0/go.mod h1:g461MYGXEXBVdV5SaR/5tNzNbSfwTBBefwc+LlDCK0Y= github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= +github.com/rivo/uniseg v0.1.0/go.mod h1:J6wj4VEh+S6ZtnVlnTBMWIodfgj8LQOQFoIToxlJtxc= +github.com/rivo/uniseg v0.2.0/go.mod h1:J6wj4VEh+S6ZtnVlnTBMWIodfgj8LQOQFoIToxlJtxc= +github.com/rivo/uniseg v0.4.7 h1:WUdvkW8uEhrYfLC4ZzdpI2ztxP1I582+49Oc5Mq64VQ= +github.com/rivo/uniseg v0.4.7/go.mod h1:FN3SvrM+Zdj16jyLfmOkMNblXMcoc8DfTHruCPUcx88= github.com/rogpeppe/go-internal v1.14.1 h1:UQB4HGPB6osV0SQTLymcB4TgvyWu6ZyliaW0tI/otEQ= github.com/rogpeppe/go-internal v1.14.1/go.mod h1:MaRKkUm5W0goXpeCfT7UZI6fk/L7L7so1lCWt35ZSgc= github.com/russross/blackfriday/v2 v2.1.0/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM= @@ -54,17 +135,37 @@ github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UV github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U= github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U= +github.com/tidwall/gjson v1.14.2/go.mod h1:/wbyibRr2FHMks5tjHJ5F8dMZh3AcwJEMf5vlfC0lxk= +github.com/tidwall/gjson v1.14.4 h1:uo0p8EbA09J7RQaflQ1aBRffTR7xedD2bcIVSYxLnkM= +github.com/tidwall/gjson v1.14.4/go.mod h1:/wbyibRr2FHMks5tjHJ5F8dMZh3AcwJEMf5vlfC0lxk= +github.com/tidwall/match v1.1.1 h1:+Ho715JplO36QYgwN9PGYNhgZvoUSc9X2c80KVTi+GA= +github.com/tidwall/match v1.1.1/go.mod h1:eRSPERbgtNPcGhD8UCthc6PmLEQXEWd3PRB5JTxsfmM= +github.com/tidwall/pretty v1.2.0/go.mod h1:ITEVvHYasfjBbM0u2Pg8T2nJnzm8xPwvNhhsoaGGjNU= +github.com/tidwall/pretty v1.2.1 h1:qjsOFOWWQl+N3RsoF5/ssm1pHmJJwhjlSbZ51I6wMl4= +github.com/tidwall/pretty v1.2.1/go.mod h1:ITEVvHYasfjBbM0u2Pg8T2nJnzm8xPwvNhhsoaGGjNU= +github.com/tidwall/sjson v1.2.5 h1:kLy8mja+1c9jlljvWTlSazM7cKDRfJuR/bOJhcY5NcY= +github.com/tidwall/sjson v1.2.5/go.mod h1:Fvgq9kS/6ociJEDnK0Fk1cpYF4FIW6ZF7LAe+6jwd28= +github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e h1:JVG44RsyaB9T2KIHavMF/ppJZNG9ZpyihvCd0w101no= +github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e/go.mod h1:RbqR21r5mrJuqunuUZ/Dhy/avygyECGrLceyNeo4LiM= github.com/yosida95/uritemplate/v3 v3.0.2 h1:Ed3Oyj9yrmi9087+NczuL5BwkIc4wvTb5zIM+UJPGz4= github.com/yosida95/uritemplate/v3 v3.0.2/go.mod h1:ILOh0sOhIJR3+L/8afwt/kE++YT040gmv5BQTMR2HP4= +github.com/yuin/goldmark v1.7.13 h1:GPddIs617DnBLFFVJFgpo1aBfe/4xcvMc3SB5t/D0pA= +github.com/yuin/goldmark v1.7.13/go.mod h1:ip/1k0VRfGynBgxOz0yCqHrbZXhcjxyuS66Brc7iBKg= +github.com/yuin/goldmark-emoji v1.0.6 h1:QWfF2FYaXwL74tfGOW5izeiZepUDroDJfWubQI9HTHs= +github.com/yuin/goldmark-emoji v1.0.6/go.mod h1:ukxJDKFpdFb5x0a5HqbdlcKtebh086iJpI31LTKmWuA= github.com/zalando/go-keyring v0.2.8 h1:6sD/Ucpl7jNq10rM2pgqTs0sZ9V3qMrqfIIy5YPccHs= github.com/zalando/go-keyring v0.2.8/go.mod h1:tsMo+VpRq5NGyKfxoBVjCuMrG47yj8cmakZDO5QGii0= go.yaml.in/yaml/v3 v3.0.4/go.mod h1:DhzuOOF2ATzADvBadXxruRBLzYTpT36CKvDb3+aBEFg= golang.org/x/crypto v0.52.0 h1:RMs7fP2rXdep0CftQlK8Uf+kibLm7qkCcradZWYz988= golang.org/x/crypto v0.52.0/go.mod h1:1QgfPxDqh0T2M/elOJtp9RvuR95kVjir0e6/BvEmGbc= +golang.org/x/net v0.54.0 h1:2zJIZAxAHV/OHCDTCOHAYehQzLfSXuf/5SoL/Dv6w/w= +golang.org/x/net v0.54.0/go.mod h1:Sj4oj8jK6XmHpBZU/zWHw3BV3abl4Kvi+Ut7cQcY+cQ= golang.org/x/oauth2 v0.36.0 h1:peZ/1z27fi9hUOFCAZaHyrpWG5lwe0RJEEEeH0ThlIs= golang.org/x/oauth2 v0.36.0/go.mod h1:YDBUJMTkDnJS+A4BP4eZBjCqtokkg1hODuPjwiGPO7Q= golang.org/x/sync v0.20.0 h1:e0PTpb7pjO8GAtTs2dQ6jYa5BWYlMuX047Dco/pItO4= golang.org/x/sync v0.20.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0= +golang.org/x/sys v0.0.0-20210809222454-d867a43fc93e/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= +golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.45.0 h1:dO4czNzziLiiXplLQgBCEpCvXQ3dnkn0SdaZSYdQ+FY= golang.org/x/sys v0.45.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw= golang.org/x/term v0.43.0 h1:S4RLU2sB31O/NCl+zFN9Aru9A/Cq2aqKpTZJ6B+DwT4= diff --git a/internal/ai/client.go b/internal/ai/client.go new file mode 100644 index 0000000..b172d1f --- /dev/null +++ b/internal/ai/client.go @@ -0,0 +1,212 @@ +package ai + +import ( + "context" + "encoding/json" + "net/http" + "strings" + + "github.com/openai/openai-go" + "github.com/openai/openai-go/option" + "github.com/openai/openai-go/shared" + + "github.com/zx06/xsql/internal/config" + "github.com/zx06/xsql/internal/errors" +) + +type ChatMessage struct { + Role string `json:"role"` + Content string `json:"content"` +} + +type Client struct { + cfg config.AIConfig + openaiClient openai.Client +} + +func NewClient(cfg config.AIConfig, httpClient *http.Client) *Client { + opts := []option.RequestOption{} + + if cfg.APIKey != "" { + opts = append(opts, option.WithAPIKey(cfg.APIKey)) + } + if cfg.BaseURL != "" { + opts = append(opts, option.WithBaseURL(strings.TrimRight(cfg.BaseURL, "/"))) + } + if httpClient != nil { + opts = append(opts, option.WithHTTPClient(httpClient)) + } + + cli := openai.NewClient(opts...) + return &Client{ + cfg: cfg, + openaiClient: cli, + } +} + +func (c *Client) ChatCompletion(ctx context.Context, messages []ChatMessage) (*AIResponse, *errors.XError) { + sdkMessages := make([]openai.ChatCompletionMessageParamUnion, 0, len(messages)) + for _, m := range messages { + switch m.Role { + case "system": + sdkMessages = append(sdkMessages, openai.SystemMessage(m.Content)) + case "user": + sdkMessages = append(sdkMessages, openai.UserMessage(m.Content)) + case "assistant": + sdkMessages = append(sdkMessages, openai.AssistantMessage(m.Content)) + default: + sdkMessages = append(sdkMessages, openai.UserMessage(m.Content)) + } + } + + sqlToolDef := openai.ChatCompletionToolParam{ + Function: shared.FunctionDefinitionParam{ + Name: "execute_sql", + Description: openai.String("Execute or present generated SQL query based on database schema and user intent"), + Parameters: shared.FunctionParameters{ + "type": "object", + "properties": map[string]interface{}{ + "sql": map[string]interface{}{ + "type": "string", + "description": "The generated SQL query statement.", + }, + "explanation": map[string]interface{}{ + "type": "string", + "description": "Explanation of what the query does or why SQL generation failed.", + }, + }, + "required": []string{"sql", "explanation"}, + }, + }, + } + + jsToolDef := openai.ChatCompletionToolParam{ + Function: shared.FunctionDefinitionParam{ + Name: "execute_javascript", + Description: openai.String("Execute JavaScript code in local goja VM sandbox to aggregate, transform, join, or format active session datasets (res1, res2, rows)."), + Parameters: shared.FunctionParameters{ + "type": "object", + "properties": map[string]interface{}{ + "js_code": map[string]interface{}{ + "type": "string", + "description": "The JavaScript code snippet to execute on active session datasets.", + }, + "explanation": map[string]interface{}{ + "type": "string", + "description": "Explanation of what the JavaScript processing code does.", + }, + }, + "required": []string{"js_code", "explanation"}, + }, + }, + } + + exportToolDef := openai.ChatCompletionToolParam{ + Function: shared.FunctionDefinitionParam{ + Name: "export_data", + Description: openai.String("Export a cached session dataset (e.g. res1, res2) to a local file in CSV, JSON, or Markdown format after human confirmation."), + Parameters: shared.FunctionParameters{ + "type": "object", + "properties": map[string]interface{}{ + "dataset_id": map[string]interface{}{ + "type": "string", + "description": "The dataset ID from session catalog to export (e.g. 'res1').", + }, + "format": map[string]interface{}{ + "type": "string", + "description": "Export file format: 'csv', 'json', or 'markdown'.", + }, + "filepath": map[string]interface{}{ + "type": "string", + "description": "Target file path (e.g. 'result.csv', 'report.json').", + }, + "explanation": map[string]interface{}{ + "type": "string", + "description": "Explanation of the data being exported.", + }, + }, + "required": []string{"dataset_id", "format", "filepath", "explanation"}, + }, + }, + } + + model := c.cfg.Model + if model == "" { + model = "gpt-4o" + } + + params := openai.ChatCompletionNewParams{ + Model: shared.ChatModel(model), + Messages: sdkMessages, + Tools: []openai.ChatCompletionToolParam{sqlToolDef, jsToolDef, exportToolDef}, + } + if c.cfg.MaxTokens > 0 { + params.MaxTokens = openai.Int(int64(c.cfg.MaxTokens)) + } + + resp, err := c.openaiClient.Chat.Completions.New(ctx, params) + if err != nil { + return nil, errors.New(errors.CodeDBExecFailed, "AI provider returned error", map[string]any{ + "err": err.Error(), + }) + } + + if len(resp.Choices) == 0 { + return nil, errors.New(errors.CodeInternal, "AI provider returned empty choices", nil) + } + + choice := resp.Choices[0] + msg := choice.Message + + for _, toolCall := range msg.ToolCalls { + switch toolCall.Function.Name { + case "execute_sql": + var raw struct { + SQL string `json:"sql"` + Explanation string `json:"explanation"` + } + if err := json.Unmarshal([]byte(toolCall.Function.Arguments), &raw); err == nil { + return &AIResponse{ + Type: TypeSQL, + SQL: strings.TrimSpace(raw.SQL), + Explanation: strings.TrimSpace(raw.Explanation), + }, nil + } + case "execute_javascript": + var raw struct { + JSCode string `json:"js_code"` + Explanation string `json:"explanation"` + } + if err := json.Unmarshal([]byte(toolCall.Function.Arguments), &raw); err == nil { + return &AIResponse{ + Type: TypeJS, + JSCode: strings.TrimSpace(raw.JSCode), + Explanation: strings.TrimSpace(raw.Explanation), + }, nil + } + case "export_data": + var raw struct { + DatasetID string `json:"dataset_id"` + Format string `json:"format"` + FilePath string `json:"filepath"` + Explanation string `json:"explanation"` + } + if err := json.Unmarshal([]byte(toolCall.Function.Arguments), &raw); err == nil { + return &AIResponse{ + Type: TypeExport, + DatasetID: strings.TrimSpace(raw.DatasetID), + Format: strings.TrimSpace(raw.Format), + FilePath: strings.TrimSpace(raw.FilePath), + Explanation: strings.TrimSpace(raw.Explanation), + }, nil + } + } + } + + content := strings.TrimSpace(msg.Content) + return &AIResponse{ + Type: TypeText, + SQL: "", + Explanation: content, + }, nil +} diff --git a/internal/ai/prompt.go b/internal/ai/prompt.go new file mode 100644 index 0000000..d13be9e --- /dev/null +++ b/internal/ai/prompt.go @@ -0,0 +1,53 @@ +package ai + +import ( + "encoding/json" + "fmt" + "strings" + + "github.com/zx06/xsql/internal/db" +) + +const SystemPromptTemplate = `You are an AI SQL Generator and Data Analyst for %s database. + +TARGET DATABASE DIALECT: %s +- Always generate correct %s SQL dialect syntax, functions, and data types. + +DATABASE SCHEMA: +%s + +%s + +ENVIRONMENT & SPECIFICATIONS: +- Database Mode: Default to READ-ONLY SELECT queries. +- JavaScript Environment: Strict ES5 (ECMAScript 5.1) engine. Active session datasets (e.g. res1, res2) are available in global context. +` + +func FormatDBName(dbType string) string { + switch strings.ToLower(dbType) { + case "mysql": + return "MySQL" + case "pg", "postgres", "postgresql": + return "PostgreSQL" + default: + if dbType != "" { + return dbType + } + return "MySQL/PostgreSQL" + } +} + +func BuildSystemPrompt(dbType string, schemaInfo *db.SchemaInfo, catalog string) string { + formattedDB := FormatDBName(dbType) + schemaJSON := "{}" + if schemaInfo != nil { + if bytes, err := json.MarshalIndent(schemaInfo, "", " "); err == nil { + schemaJSON = string(bytes) + } + } + catalogBlock := "" + if catalog != "" { + catalogBlock = fmt.Sprintf("SESSION DATASETS CATALOG:\n%s\n", catalog) + } + return fmt.Sprintf(SystemPromptTemplate, formattedDB, formattedDB, formattedDB, schemaJSON, catalogBlock) +} diff --git a/internal/ai/service.go b/internal/ai/service.go new file mode 100644 index 0000000..e6baa9a --- /dev/null +++ b/internal/ai/service.go @@ -0,0 +1,62 @@ +package ai + +import ( + "context" + + "github.com/zx06/xsql/internal/config" + "github.com/zx06/xsql/internal/db" + "github.com/zx06/xsql/internal/errors" +) + +type ResponseType string + +const ( + TypeSQL ResponseType = "sql" + TypeJS ResponseType = "js" + TypeExport ResponseType = "export" + TypeText ResponseType = "text" +) + +type AIResponse struct { + Type ResponseType `json:"type"` + SQL string `json:"sql,omitempty"` + JSCode string `json:"js_code,omitempty"` + DatasetID string `json:"dataset_id,omitempty"` + Format string `json:"format,omitempty"` + FilePath string `json:"filepath,omitempty"` + Explanation string `json:"explanation"` +} + +type SQLResponse = AIResponse + +type Service struct { + client *Client +} + +func NewService(cfg config.AIConfig, client *Client) *Service { + if client == nil { + client = NewClient(cfg, nil) + } + return &Service{ + client: client, + } +} + +func (s *Service) GenerateSQL(ctx context.Context, userPrompt string, schemaInfo *db.SchemaInfo, dbType string) (*AIResponse, *errors.XError) { + return s.GenerateResponse(ctx, userPrompt, schemaInfo, dbType, "") +} + +func (s *Service) GenerateResponse(ctx context.Context, userPrompt string, schemaInfo *db.SchemaInfo, dbType string, catalog string) (*AIResponse, *errors.XError) { + systemPrompt := BuildSystemPrompt(dbType, schemaInfo, catalog) + + messages := []ChatMessage{ + {Role: "system", Content: systemPrompt}, + {Role: "user", Content: userPrompt}, + } + + return s.client.ChatCompletion(ctx, messages) +} + +func (s *Service) ChatCompletion(ctx context.Context, messages []ChatMessage) (*AIResponse, *errors.XError) { + return s.client.ChatCompletion(ctx, messages) +} diff --git a/internal/ai/service_test.go b/internal/ai/service_test.go new file mode 100644 index 0000000..f5bd174 --- /dev/null +++ b/internal/ai/service_test.go @@ -0,0 +1,279 @@ +package ai + +import ( + "context" + "net/http" + "net/http/httptest" + "testing" + + "github.com/zx06/xsql/internal/config" + "github.com/zx06/xsql/internal/db" +) + +func TestBuildSystemPrompt(t *testing.T) { + schema := &db.SchemaInfo{ + Database: "testdb", + Tables: []db.Table{ + { + Name: "users", + Columns: []db.Column{ + {Name: "id", Type: "int", PrimaryKey: true}, + {Name: "name", Type: "varchar"}, + }, + }, + }, + } + + prompt := BuildSystemPrompt("mysql", schema, "res1: users") + if prompt == "" { + t.Fatal("expected non-empty prompt") + } + + defaultPrompt := BuildSystemPrompt("", nil, "") + if defaultPrompt == "" { + t.Fatal("expected non-empty default prompt") + } +} + +func TestNewService_NilClient(t *testing.T) { + cfg := config.AIConfig{ + Provider: "openai", + APIKey: "key", + } + svc := NewService(cfg, nil) + if svc == nil || svc.client == nil { + t.Fatal("expected non-nil Service and Client") + } +} + +func TestGenerateSQL_MockHTTP_ToolCall(t *testing.T) { + mockServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/chat/completions" { + t.Errorf("unexpected path: %s", r.URL.Path) + } + if r.Header.Get("Authorization") != "Bearer test-key" { + t.Errorf("unexpected auth header: %s", r.Header.Get("Authorization")) + } + + respBody := `{ + "id": "chatcmpl-123", + "object": "chat.completion", + "created": 1677652288, + "model": "gpt-4o", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": null, + "tool_calls": [ + { + "id": "call_abc123", + "type": "function", + "function": { + "name": "execute_sql", + "arguments": "{\"sql\":\"SELECT id, name FROM users;\",\"explanation\":\"Queries all users.\"}" + } + } + ] + }, + "finish_reason": "tool_calls" + } + ] + }` + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(respBody)) + })) + defer mockServer.Close() + + cfg := config.AIConfig{ + Provider: "openai", + BaseURL: mockServer.URL, + APIKey: "test-key", + Model: "gpt-4o", + MaxTokens: 100, + } + + client := NewClient(cfg, mockServer.Client()) + service := NewService(cfg, client) + + res, xe := service.GenerateSQL(context.Background(), "show users", nil, "mysql") + if xe != nil { + t.Fatalf("unexpected error: %v", xe) + } + + if res.SQL != "SELECT id, name FROM users;" { + t.Errorf("expected SQL 'SELECT id, name FROM users;', got %q", res.SQL) + } + if res.Explanation != "Queries all users." { + t.Errorf("expected explanation 'Queries all users.', got %q", res.Explanation) + } +} + +func TestGenerateResponse_MockHTTP_JSToolCall(t *testing.T) { + mockServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + respBody := `{ + "id": "chatcmpl-125", + "object": "chat.completion", + "created": 1677652288, + "model": "gpt-4o", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": null, + "tool_calls": [ + { + "id": "call_js123", + "type": "function", + "function": { + "name": "execute_javascript", + "arguments": "{\"js_code\":\"rows.filter(r => r.status === 'ONLINE');\",\"explanation\":\"Filters online servers.\"}" + } + } + ] + }, + "finish_reason": "tool_calls" + } + ] + }` + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(respBody)) + })) + defer mockServer.Close() + + cfg := config.AIConfig{ + Provider: "openai", + BaseURL: mockServer.URL, + APIKey: "test-key", + } + + client := NewClient(cfg, mockServer.Client()) + service := NewService(cfg, client) + + res, xe := service.GenerateResponse(context.Background(), "filter online servers", nil, "mysql", "res1 catalog") + if xe != nil { + t.Fatalf("unexpected error: %v", xe) + } + + if res.Type != TypeJS { + t.Errorf("expected type JS, got %q", res.Type) + } + if res.JSCode != "rows.filter(r => r.status === 'ONLINE');" { + t.Errorf("unexpected JS code: %q", res.JSCode) + } +} + +func TestGenerateSQL_MockHTTP_TextMessageFallback(t *testing.T) { + mockServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + respBody := `{ + "id": "chatcmpl-124", + "object": "chat.completion", + "created": 1677652288, + "model": "gpt-4o", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "I cannot answer this question based on the schema." + }, + "finish_reason": "stop" + } + ] + }` + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(respBody)) + })) + defer mockServer.Close() + + cfg := config.AIConfig{ + Provider: "openai", + BaseURL: mockServer.URL, + APIKey: "test-key", + Model: "gpt-4o", + } + + client := NewClient(cfg, mockServer.Client()) + service := NewService(cfg, client) + + res, xe := service.GenerateSQL(context.Background(), "unknown table", nil, "mysql") + if xe != nil { + t.Fatalf("unexpected error: %v", xe) + } + + if res.SQL != "" { + t.Errorf("expected empty SQL, got %q", res.SQL) + } + if res.Explanation != "I cannot answer this question based on the schema." { + t.Errorf("expected explanation, got %q", res.Explanation) + } +} + +func TestGenerateSQL_MockHTTP_APIError(t *testing.T) { + mockServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusInternalServerError) + _, _ = w.Write([]byte(`{"error": {"message": "internal server error"}}`)) + })) + defer mockServer.Close() + + cfg := config.AIConfig{ + Provider: "openai", + BaseURL: mockServer.URL, + APIKey: "test-key", + } + + client := NewClient(cfg, mockServer.Client()) + service := NewService(cfg, client) + + _, xe := service.GenerateSQL(context.Background(), "test", nil, "mysql") + if xe == nil { + t.Fatal("expected error for HTTP 500") + } +} + +func TestService_ChatCompletion(t *testing.T) { + mockServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + respBody := `{ + "id": "chatcmpl-999", + "object": "chat.completion", + "created": 1677652288, + "model": "gpt-4o", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "Hello from ChatCompletion" + }, + "finish_reason": "stop" + } + ] + }` + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(respBody)) + })) + defer mockServer.Close() + + cfg := config.AIConfig{ + Provider: "openai", + BaseURL: mockServer.URL, + APIKey: "test-key", + } + + client := NewClient(cfg, mockServer.Client()) + service := NewService(cfg, client) + + msgs := []ChatMessage{ + {Role: "user", Content: "Hello"}, + } + + res, xe := service.ChatCompletion(context.Background(), msgs) + if xe != nil { + t.Fatalf("unexpected error: %v", xe) + } + if res.Explanation != "Hello from ChatCompletion" { + t.Errorf("expected explanation 'Hello from ChatCompletion', got %q", res.Explanation) + } +} diff --git a/internal/app/conn.go b/internal/app/conn.go index 9e1496c..c76b480 100644 --- a/internal/app/conn.go +++ b/internal/app/conn.go @@ -7,6 +7,8 @@ import ( "github.com/zx06/xsql/internal/config" "github.com/zx06/xsql/internal/db" + _ "github.com/zx06/xsql/internal/db/mysql" + _ "github.com/zx06/xsql/internal/db/pg" "github.com/zx06/xsql/internal/errors" "github.com/zx06/xsql/internal/secret" "github.com/zx06/xsql/internal/ssh" diff --git a/internal/config/resolve.go b/internal/config/resolve.go index b9e4f1e..d1bfa78 100644 --- a/internal/config/resolve.go +++ b/internal/config/resolve.go @@ -61,33 +61,39 @@ func Resolve(opts Options) (Resolved, *errors.XError) { } } + // Resolve all profiles in cfg.Profiles + resolvedProfiles := make(map[string]Profile, len(cfg.Profiles)) + for name, p := range cfg.Profiles { + pCopy := p + if pCopy.SSHProxy != "" { + if proxy, ok := cfg.SSHProxies[pCopy.SSHProxy]; ok { + pCopy.SSHConfig = &proxy + } + } + if pCopy.Port == 0 { + switch pCopy.DB { + case "mysql": + pCopy.Port = 3306 + case "pg": + pCopy.Port = 5432 + } + } + resolvedProfiles[name] = pCopy + } + // 3) Retrieve full profile var selectedProfile Profile if profile != "" { - p, ok := cfg.Profiles[profile] + p, ok := resolvedProfiles[profile] if !ok { return Resolved{}, errors.New(errors.CodeCfgInvalid, "profile not found", map[string]any{"profile": profile}) } - selectedProfile = p - // Resolve ssh_proxy reference - if selectedProfile.SSHProxy != "" { - if proxy, ok := cfg.SSHProxies[selectedProfile.SSHProxy]; ok { - selectedProfile.SSHConfig = &proxy - } else { - return Resolved{}, errors.New(errors.CodeCfgInvalid, "ssh_proxy not found", - map[string]any{"profile": profile, "ssh_proxy": selectedProfile.SSHProxy}) - } - } - // Set default port - if selectedProfile.Port == 0 { - switch selectedProfile.DB { - case "mysql": - selectedProfile.Port = 3306 - case "pg": - selectedProfile.Port = 5432 - } + if p.SSHProxy != "" && p.SSHConfig == nil { + return Resolved{}, errors.New(errors.CodeCfgInvalid, "ssh_proxy not found", + map[string]any{"profile": profile, "ssh_proxy": p.SSHProxy}) } + selectedProfile = p } // 4) Merge format: --format > XSQL_FORMAT > profile.format > auto @@ -102,5 +108,47 @@ func Resolve(opts Options) (Resolved, *errors.XError) { format = opts.CLIFormat } - return Resolved{ConfigPath: cfgPath, ProfileName: profile, Format: format, Profile: selectedProfile}, nil + // 5) Merge AI Config: CLI > ENV > Config > Default + aiConfig := cfg.AI + if aiConfig.Provider == "" { + aiConfig.Provider = "openai" + } + if aiConfig.BaseURL == "" { + aiConfig.BaseURL = "https://api.openai.com/v1" + } + if aiConfig.Model == "" { + aiConfig.Model = "gpt-4o" + } + if aiConfig.MaxTokens == 0 { + aiConfig.MaxTokens = 2048 + } + + if opts.EnvAIBaseURL != "" { + aiConfig.BaseURL = opts.EnvAIBaseURL + } + if opts.EnvAIModel != "" { + aiConfig.Model = opts.EnvAIModel + } + if opts.EnvAIAPIKey != "" { + aiConfig.APIKey = opts.EnvAIAPIKey + } + + if opts.CLIAIBaseURLSet && opts.CLIAIBaseURL != "" { + aiConfig.BaseURL = opts.CLIAIBaseURL + } + if opts.CLIAIModelSet && opts.CLIAIModel != "" { + aiConfig.Model = opts.CLIAIModel + } + if opts.CLIAIAPIKeySet && opts.CLIAIAPIKey != "" { + aiConfig.APIKey = opts.CLIAIAPIKey + } + + return Resolved{ + ConfigPath: cfgPath, + ProfileName: profile, + Format: format, + Profile: selectedProfile, + AllProfiles: resolvedProfiles, + AI: aiConfig, + }, nil } diff --git a/internal/config/resolve_test.go b/internal/config/resolve_test.go index ceee2ba..a281b63 100644 --- a/internal/config/resolve_test.go +++ b/internal/config/resolve_test.go @@ -216,3 +216,61 @@ func TestResolve_Port_PreservedWhenSpecified(t *testing.T) { t.Errorf("expected port 13306 to be preserved, got %d", got.Profile.Port) } } + +func TestResolve_AIConfigPrecedence(t *testing.T) { + tmp := t.TempDir() + cfg := []byte(`ai: + base_url: "https://config.api.com" + model: "config-model" + api_key: "config-key" +`) + path := filepath.Join(tmp, "xsql.yaml") + if err := os.WriteFile(path, cfg, 0o600); err != nil { + t.Fatal(err) + } + + // 1. Config values + got, xe := Resolve(Options{WorkDir: tmp, HomeDir: tmp}) + if xe != nil { + t.Fatal(xe) + } + if got.AI.BaseURL != "https://config.api.com" || got.AI.Model != "config-model" || got.AI.APIKey != "config-key" { + t.Fatalf("unexpected AI config from yaml: %+v", got.AI) + } + + // 2. ENV overrides config + got, xe = Resolve(Options{ + WorkDir: tmp, + HomeDir: tmp, + EnvAIBaseURL: "https://env.api.com", + EnvAIModel: "env-model", + EnvAIAPIKey: "env-key", + }) + if xe != nil { + t.Fatal(xe) + } + if got.AI.BaseURL != "https://env.api.com" || got.AI.Model != "env-model" || got.AI.APIKey != "env-key" { + t.Fatalf("unexpected AI config from env: %+v", got.AI) + } + + // 3. CLI overrides ENV & Config + got, xe = Resolve(Options{ + WorkDir: tmp, + HomeDir: tmp, + EnvAIBaseURL: "https://env.api.com", + EnvAIModel: "env-model", + EnvAIAPIKey: "env-key", + CLIAIBaseURL: "https://cli.api.com", + CLIAIBaseURLSet: true, + CLIAIModel: "cli-model", + CLIAIModelSet: true, + CLIAIAPIKey: "cli-key", + CLIAIAPIKeySet: true, + }) + if xe != nil { + t.Fatal(xe) + } + if got.AI.BaseURL != "https://cli.api.com" || got.AI.Model != "cli-model" || got.AI.APIKey != "cli-key" { + t.Fatalf("unexpected AI config from cli: %+v", got.AI) + } +} diff --git a/internal/config/types.go b/internal/config/types.go index a07d5db..88b2455 100644 --- a/internal/config/types.go +++ b/internal/config/types.go @@ -10,6 +10,16 @@ type File struct { MCP MCPConfig `yaml:"mcp" json:"mcp"` Web WebConfig `yaml:"web" json:"web"` Stats stats.StatsConfig `yaml:"stats" json:"stats"` + AI AIConfig `yaml:"ai" json:"ai"` +} + +// AIConfig defines the AI LLM service configuration. +type AIConfig struct { + Provider string `yaml:"provider" json:"provider"` // default "openai" + BaseURL string `yaml:"base_url" json:"base_url"` // default "https://api.openai.com/v1" + APIKey string `yaml:"api_key" json:"api_key"` // supports keyring:xxx reference + Model string `yaml:"model" json:"model"` // default "gpt-4o" + MaxTokens int `yaml:"max_tokens" json:"max_tokens"` // default 2048 } // SSHProxy defines a reusable SSH proxy configuration. @@ -83,7 +93,9 @@ type Resolved struct { ConfigPath string ProfileName string Format string - Profile Profile // full profile for query use + Profile Profile // full profile for query use + AllProfiles map[string]Profile // all configured profiles + AI AIConfig } type Options struct { @@ -96,9 +108,20 @@ type Options struct { CLIFormat string CLIFormatSet bool + // CLI AI + CLIAIModel string + CLIAIModelSet bool + CLIAIBaseURL string + CLIAIBaseURLSet bool + CLIAIAPIKey string + CLIAIAPIKeySet bool + // ENV (injected by caller for testability) - EnvProfile string - EnvFormat string + EnvProfile string + EnvFormat string + EnvAIModel string + EnvAIBaseURL string + EnvAIAPIKey string // HomeDir is used for default path resolution (auto-detected if empty). HomeDir string diff --git a/internal/export/exporter.go b/internal/export/exporter.go new file mode 100644 index 0000000..7822bb2 --- /dev/null +++ b/internal/export/exporter.go @@ -0,0 +1,117 @@ +package export + +import ( + "encoding/csv" + "encoding/json" + "fmt" + "os" + "path/filepath" + "strings" + + "github.com/zx06/xsql/internal/db" + "github.com/zx06/xsql/internal/errors" +) + +type ExportFormat string + +const ( + FormatCSV ExportFormat = "csv" + FormatJSON ExportFormat = "json" + FormatMarkdown ExportFormat = "markdown" +) + +func ExportQueryResult(result *db.QueryResult, format ExportFormat, filePath string) (string, *errors.XError) { + if result == nil { + return "", errors.New(errors.CodeCfgInvalid, "cannot export nil QueryResult", nil) + } + + if filePath == "" { + filePath = fmt.Sprintf("export_%s.%s", format, format) + } + + // Ensure directory exists + dir := filepath.Dir(filePath) + if dir != "" && dir != "." { + if err := os.MkdirAll(dir, 0755); err != nil { + return "", errors.New(errors.CodeInternal, "failed to create export directory", map[string]any{ + "dir": dir, + "err": err.Error(), + }) + } + } + + f, err := os.Create(filePath) + if err != nil { + return "", errors.New(errors.CodeInternal, "failed to create export file", map[string]any{ + "path": filePath, + "err": err.Error(), + }) + } + defer func() { _ = f.Close() }() + + switch format { + case FormatJSON: + enc := json.NewEncoder(f) + enc.SetIndent("", " ") + if err := enc.Encode(result.Rows); err != nil { + return "", errors.New(errors.CodeInternal, "failed to write JSON export", map[string]any{"err": err.Error()}) + } + + case FormatMarkdown: + var sb strings.Builder + sb.WriteString("| " + strings.Join(result.Columns, " | ") + " |\n") + var sep []string + for range result.Columns { + sep = append(sep, "---") + } + sb.WriteString("| " + strings.Join(sep, " | ") + " |\n") + + for _, row := range result.Rows { + var vals []string + for _, col := range result.Columns { + val := row[col] + if val == nil { + vals = append(vals, "NULL") + } else { + cellStr := fmt.Sprintf("%v", val) + cellStr = strings.ReplaceAll(cellStr, "\n", " ") + cellStr = strings.ReplaceAll(cellStr, "|", "\\|") + vals = append(vals, cellStr) + } + } + sb.WriteString("| " + strings.Join(vals, " | ") + " |\n") + } + if _, err := f.WriteString(sb.String()); err != nil { + return "", errors.New(errors.CodeInternal, "failed to write Markdown export", map[string]any{"err": err.Error()}) + } + + case FormatCSV: + fallthrough + default: + w := csv.NewWriter(f) + if err := w.Write(result.Columns); err != nil { + return "", errors.New(errors.CodeInternal, "failed to write CSV header", map[string]any{"err": err.Error()}) + } + for _, row := range result.Rows { + var vals []string + for _, col := range result.Columns { + val := row[col] + if val == nil { + vals = append(vals, "") + } else { + vals = append(vals, fmt.Sprintf("%v", val)) + } + } + if err := w.Write(vals); err != nil { + return "", errors.New(errors.CodeInternal, "failed to write CSV row", map[string]any{"err": err.Error()}) + } + } + w.Flush() + if err := w.Error(); err != nil { + return "", errors.New(errors.CodeInternal, "failed to flush CSV writer", map[string]any{"err": err.Error()}) + } + } + + absPath, _ := filepath.Abs(filePath) + return absPath, nil +} diff --git a/internal/export/exporter_test.go b/internal/export/exporter_test.go new file mode 100644 index 0000000..ff53f18 --- /dev/null +++ b/internal/export/exporter_test.go @@ -0,0 +1,72 @@ +package export + +import ( + "os" + "path/filepath" + "strings" + "testing" + + "github.com/zx06/xsql/internal/db" +) + +func TestExportQueryResult_CSV_JSON_MD(t *testing.T) { + tempDir, err := os.MkdirTemp("", "xsql_export_test") + if err != nil { + t.Fatalf("failed to create temp dir: %v", err) + } + defer os.RemoveAll(tempDir) + + res := &db.QueryResult{ + Columns: []string{"id", "username", "status"}, + Rows: []map[string]any{ + {"id": 1, "username": "alice", "status": "active"}, + {"id": 2, "username": "bob", "status": nil}, + }, + } + + // 1. CSV + csvPath := filepath.Join(tempDir, "test.csv") + absPath, xe := ExportQueryResult(res, FormatCSV, csvPath) + if xe != nil { + t.Fatalf("CSV export failed: %v", xe) + } + content, _ := os.ReadFile(absPath) + if !strings.Contains(string(content), "username") || !strings.Contains(string(content), "alice") { + t.Fatalf("unexpected CSV content: %s", string(content)) + } + + // 2. JSON + jsonPath := filepath.Join(tempDir, "test.json") + absPath, xe = ExportQueryResult(res, FormatJSON, jsonPath) + if xe != nil { + t.Fatalf("JSON export failed: %v", xe) + } + content, _ = os.ReadFile(absPath) + if !strings.Contains(string(content), `"alice"`) { + t.Fatalf("unexpected JSON content: %s", string(content)) + } + + // 3. Markdown + mdPath := filepath.Join(tempDir, "sub", "test.md") + res.Rows[1]["status"] = "multiline\ntext|pipe" + absPath, xe = ExportQueryResult(res, FormatMarkdown, mdPath) + if xe != nil { + t.Fatalf("Markdown export failed: %v", xe) + } + content, _ = os.ReadFile(absPath) + if !strings.Contains(string(content), "| username |") || !strings.Contains(string(content), "text\\|pipe") { + t.Fatalf("unexpected Markdown content: %s", string(content)) + } + + // 4. Nil result & empty filePath fallback + _, xe = ExportQueryResult(nil, FormatCSV, "") + if xe == nil { + t.Fatal("expected error for nil QueryResult") + } + + absDefault, xe := ExportQueryResult(res, FormatCSV, "") + if xe != nil { + t.Fatalf("unexpected error for empty filePath: %v", xe) + } + _ = os.Remove(absDefault) +} diff --git a/internal/js/engine.go b/internal/js/engine.go new file mode 100644 index 0000000..a39b34f --- /dev/null +++ b/internal/js/engine.go @@ -0,0 +1,158 @@ +package js + +import ( + "context" + "encoding/json" + "fmt" + "strings" + "time" + + "github.com/dop251/goja" + + "github.com/zx06/xsql/internal/errors" + "github.com/zx06/xsql/internal/session" +) + +type ExecutionResult struct { + Value any `json:"value"` + JSONString string `json:"json_string"` + SummaryText string `json:"summary_text"` + Logs []string `json:"logs"` +} + +type JSEngine struct { + DefaultTimeout time.Duration +} + +func NewJSEngine(timeout time.Duration) *JSEngine { + if timeout <= 0 { + timeout = 1 * time.Minute + } + return &JSEngine{ + DefaultTimeout: timeout, + } +} + +func (e *JSEngine) Execute(ctx context.Context, jsCode string, store *session.SessionDataStore) (*ExecutionResult, *errors.XError) { + if ctx == nil { + ctx = context.Background() + } + + execCtx, cancel := context.WithTimeout(ctx, e.DefaultTimeout) + defer cancel() + + vm := goja.New() + vm.SetFieldNameMapper(goja.TagFieldNameMapper("json", true)) + + // Inject console.log / console.error capture + var logs []string + console := vm.NewObject() + _ = console.Set("log", func(call goja.FunctionCall) goja.Value { + var args []string + for _, arg := range call.Arguments { + args = append(args, fmt.Sprintf("%v", arg.Export())) + } + logs = append(logs, strings.Join(args, " ")) + return goja.Undefined() + }) + _ = console.Set("error", func(call goja.FunctionCall) goja.Value { + var args []string + for _, arg := range call.Arguments { + args = append(args, fmt.Sprintf("%v", arg.Export())) + } + logs = append(logs, "[ERROR] "+strings.Join(args, " ")) + return goja.Undefined() + }) + _ = vm.Set("console", console) + + // Inject all active datasets from store + if store != nil { + allData := store.GetAll() + for id, queryRes := range allData { + if queryRes != nil { + _ = vm.Set(id, queryRes.Rows) + } + } + // Inject latest query result as `rows` and `columns` + if latest, ok := store.Latest(); ok && latest != nil { + _ = vm.Set("rows", latest.Rows) + _ = vm.Set("columns", latest.Columns) + } + } + + // Timeout interrupt setup + doneChan := make(chan struct{}) + defer close(doneChan) + + go func() { + select { + case <-execCtx.Done(): + if execCtx.Err() == context.DeadlineExceeded { + vm.Interrupt("execution timeout (limit reached)") + } + case <-doneChan: + } + }() + + // Execute JS script + val, err := vm.RunString(jsCode) + if err != nil { + return nil, errors.New(errors.CodeDBExecFailed, fmt.Sprintf("JavaScript execution error: %v", err), map[string]any{ + "err": err.Error(), + }) + } + + if val == nil || goja.IsUndefined(val) || goja.IsNull(val) { + summary := "(null)" + if len(logs) > 0 { + summary = strings.Join(logs, "\n") + } + return &ExecutionResult{ + Value: nil, + JSONString: "null", + SummaryText: summary, + Logs: logs, + }, nil + } + + exported := val.Export() + var jsonStr string + var summaryText string + + if str, ok := exported.(string); ok { + var parsed any + if json.Unmarshal([]byte(str), &parsed) == nil { + if b, err := json.MarshalIndent(parsed, "", " "); err == nil { + jsonStr = string(b) + } else { + jsonStr = str + } + } else { + jsonStr = str + } + summaryText = jsonStr + } else if exported != nil { + if b, err := json.MarshalIndent(exported, "", " "); err == nil { + jsonStr = string(b) + } else { + jsonStr = fmt.Sprintf("%v", exported) + } + summaryText = jsonStr + } + + if len(logs) > 0 { + consoleLogs := strings.Join(logs, "\n") + if summaryText == "" || summaryText == "(null)" { + summaryText = consoleLogs + } else { + summaryText = consoleLogs + "\n\n" + summaryText + } + } + + return &ExecutionResult{ + Value: exported, + JSONString: jsonStr, + SummaryText: summaryText, + Logs: logs, + }, nil +} diff --git a/internal/js/engine_test.go b/internal/js/engine_test.go new file mode 100644 index 0000000..8710d3f --- /dev/null +++ b/internal/js/engine_test.go @@ -0,0 +1,149 @@ +package js + +import ( + "context" + "strings" + "testing" + "time" + + "github.com/zx06/xsql/internal/db" + "github.com/zx06/xsql/internal/session" +) + +func TestJSEngine_MultiDatasetRecall(t *testing.T) { + store := session.NewSessionDataStore() + + res1 := &db.QueryResult{ + Columns: []string{"id", "name", "status"}, + Rows: []map[string]any{ + {"id": 1, "name": "srv1", "status": "ONLINE"}, + {"id": 2, "name": "srv2", "status": "OFFLINE"}, + }, + } + store.Save("servers query", res1) + + res2 := &db.QueryResult{ + Columns: []string{"id", "server_id", "severity"}, + Rows: []map[string]any{ + {"id": 101, "server_id": 1, "severity": "HIGH"}, + }, + } + store.Save("alerts query", res2) + + engine := NewJSEngine(5 * time.Second) + + // JS script joining res1 and res2 + jsCode := ` + (function() { + var alertServerIds = new Set(res2.map(function(a) { return a.server_id; })); + var onlineWithAlerts = res1.filter(function(s) { + return s.status === 'ONLINE' && alertServerIds.has(s.id); + }); + return { + count: onlineWithAlerts.length, + servers: onlineWithAlerts + }; + })(); + ` + + result, xe := engine.Execute(context.Background(), jsCode, store) + if xe != nil { + t.Fatalf("unexpected execution error: %v", xe) + } + + if !strings.Contains(result.JSONString, `"count": 1`) || !strings.Contains(result.JSONString, `"srv1"`) { + t.Fatalf("expected joined result with srv1, got:\n%s", result.JSONString) + } +} + +func TestJSEngine_Timeout(t *testing.T) { + engine := NewJSEngine(100 * time.Millisecond) + + jsCode := ` + while(true) {} + ` + + _, xe := engine.Execute(context.Background(), jsCode, nil) + if xe == nil { + t.Fatal("expected timeout error for infinite loop, got nil") + } +} + +func TestJSEngine_SyntaxAndRuntimeError(t *testing.T) { + engine := NewJSEngine(1 * time.Second) + + _, xe := engine.Execute(context.Background(), "invalid js {code", nil) + if xe == nil { + t.Fatal("expected syntax error for invalid JS") + } + + _, xe = engine.Execute(context.Background(), "throw new Error('custom js error');", nil) + if xe == nil { + t.Fatal("expected runtime error for thrown error") + } +} + +func TestJSEngine_PrimitiveResult(t *testing.T) { + engine := NewJSEngine(1 * time.Second) + + res, xe := engine.Execute(context.Background(), "'hello world'", nil) + if xe != nil { + t.Fatalf("unexpected error: %v", xe) + } + if res.SummaryText != "hello world" { + t.Fatalf("expected summary 'hello world', got %q", res.SummaryText) + } + + res, xe = engine.Execute(context.Background(), "[1, 2, 3]", nil) + if xe != nil { + t.Fatalf("unexpected error: %v", xe) + } + if !strings.Contains(res.JSONString, "[1, 2, 3]") && !strings.Contains(res.JSONString, "[\n 1") { + t.Fatalf("expected array json output, got %q", res.JSONString) + } +} + +func TestJSEngine_ConsoleAndNullAndDefaultTimeout(t *testing.T) { + engine := NewJSEngine(0) // Default 1 min + if engine.DefaultTimeout <= 0 { + t.Fatal("expected positive default timeout") + } + + store := session.NewSessionDataStore() + store.Save("test query", &db.QueryResult{ + Columns: []string{"id"}, + Rows: []map[string]any{{"id": 10}}, + }) + + // 1. Console log and error with undefined return + jsCode := ` + console.log("log msg", 123); + console.error("error msg"); + rows.length; + ` + res, xe := engine.Execute(nil, jsCode, store) + if xe != nil { + t.Fatalf("unexpected execution error: %v", xe) + } + if len(res.Logs) != 2 || !strings.Contains(res.SummaryText, "[ERROR] error msg") { + t.Fatalf("expected 2 log entries in summary, got:\n%s", res.SummaryText) + } + + // 2. Null/Undefined return with console logs + resNull, xe := engine.Execute(nil, `console.log("hello"); null;`, nil) + if xe != nil { + t.Fatalf("unexpected execution error: %v", xe) + } + if resNull.JSONString != "null" || !strings.Contains(resNull.SummaryText, "hello") { + t.Fatalf("expected null return with logs, got %q", resNull.SummaryText) + } + + // 3. JSON String return + resJSON, xe := engine.Execute(nil, `JSON.stringify({status: "ok"})`, nil) + if xe != nil { + t.Fatalf("unexpected execution error: %v", xe) + } + if !strings.Contains(resJSON.JSONString, `"status": "ok"`) { + t.Fatalf("expected pretty formatted JSON string, got %q", resJSON.JSONString) + } +} diff --git a/internal/session/store.go b/internal/session/store.go new file mode 100644 index 0000000..fbafeaa --- /dev/null +++ b/internal/session/store.go @@ -0,0 +1,108 @@ +package session + +import ( + "fmt" + "strings" + "sync" + + "github.com/zx06/xsql/internal/db" +) + +type DatasetEntry struct { + ID string `json:"id"` + Description string `json:"description"` + Result *db.QueryResult `json:"result"` +} + +type SessionDataStore struct { + mu sync.RWMutex + datasets map[string]*DatasetEntry + order []string + counter int +} + +func NewSessionDataStore() *SessionDataStore { + return &SessionDataStore{ + datasets: make(map[string]*DatasetEntry), + order: make([]string, 0), + } +} + +func (s *SessionDataStore) Save(description string, result *db.QueryResult) string { + if result == nil { + return "" + } + s.mu.Lock() + defer s.mu.Unlock() + + s.counter++ + id := fmt.Sprintf("res%d", s.counter) + + entry := &DatasetEntry{ + ID: id, + Description: strings.TrimSpace(description), + Result: result, + } + + s.datasets[id] = entry + s.order = append(s.order, id) + return id +} + +func (s *SessionDataStore) Get(id string) (*db.QueryResult, bool) { + s.mu.RLock() + defer s.mu.RUnlock() + + entry, ok := s.datasets[id] + if !ok || entry == nil { + return nil, false + } + return entry.Result, true +} + +func (s *SessionDataStore) Latest() (*db.QueryResult, bool) { + s.mu.RLock() + defer s.mu.RUnlock() + + if len(s.order) == 0 { + return nil, false + } + lastID := s.order[len(s.order)-1] + return s.datasets[lastID].Result, true +} + +func (s *SessionDataStore) GetAll() map[string]*db.QueryResult { + s.mu.RLock() + defer s.mu.RUnlock() + + res := make(map[string]*db.QueryResult, len(s.datasets)) + for k, v := range s.datasets { + res[k] = v.Result + } + return res +} + +func (s *SessionDataStore) GetCatalog() string { + s.mu.RLock() + defer s.mu.RUnlock() + + if len(s.datasets) == 0 { + return "(No active datasets in current session)" + } + + var sb strings.Builder + sb.WriteString("Active Datasets in Session:\n") + for _, id := range s.order { + entry := s.datasets[id] + if entry == nil || entry.Result == nil { + continue + } + cols := strings.Join(entry.Result.Columns, ", ") + desc := entry.Description + if desc == "" { + desc = "Query Result" + } + fmt.Fprintf(&sb, "- `%s`: %s (%d rows, columns: [%s])\n", id, desc, len(entry.Result.Rows), cols) + } + return sb.String() +} diff --git a/internal/session/store_test.go b/internal/session/store_test.go new file mode 100644 index 0000000..eeef150 --- /dev/null +++ b/internal/session/store_test.go @@ -0,0 +1,54 @@ +package session + +import ( + "strings" + "testing" + + "github.com/zx06/xsql/internal/db" +) + +func TestSessionDataStore(t *testing.T) { + store := NewSessionDataStore() + + if catalog := store.GetCatalog(); !strings.Contains(catalog, "No active datasets") { + t.Fatalf("expected empty catalog warning, got: %s", catalog) + } + + res1 := &db.QueryResult{ + Columns: []string{"id", "name"}, + Rows: []map[string]any{{"id": 1, "name": "srv1"}}, + } + id1 := store.Save("servers query", res1) + if id1 != "res1" { + t.Fatalf("expected ID 'res1', got %s", id1) + } + + res2 := &db.QueryResult{ + Columns: []string{"id", "severity"}, + Rows: []map[string]any{{"id": 1, "severity": "HIGH"}}, + } + id2 := store.Save("alerts query", res2) + if id2 != "res2" { + t.Fatalf("expected ID 'res2', got %s", id2) + } + + latest, ok := store.Latest() + if !ok || len(latest.Rows) != 1 || latest.Rows[0]["severity"] != "HIGH" { + t.Fatal("failed to retrieve latest dataset") + } + + catalog := store.GetCatalog() + if !strings.Contains(catalog, "`res1`") || !strings.Contains(catalog, "`res2`") { + t.Fatalf("expected catalog to list res1 and res2, got:\n%s", catalog) + } + + got1, ok := store.Get("res1") + if !ok || got1.Rows[0]["name"] != "srv1" { + t.Fatal("failed to get res1 from store") + } + + all := store.GetAll() + if len(all) != 2 || all["res1"] == nil || all["res2"] == nil { + t.Fatalf("expected GetAll to return 2 datasets, got %d", len(all)) + } +} diff --git a/internal/tui/components.go b/internal/tui/components.go new file mode 100644 index 0000000..d12f1f0 --- /dev/null +++ b/internal/tui/components.go @@ -0,0 +1,307 @@ +package tui + +import ( + "fmt" + "strings" + + "github.com/charmbracelet/lipgloss" + "github.com/charmbracelet/lipgloss/table" + "github.com/mattn/go-runewidth" + + "github.com/zx06/xsql/internal/db" +) + +var ( + TableHeaderStyle = lipgloss.NewStyle(). + Bold(true). + Foreground(lipgloss.AdaptiveColor{Light: "#4F46E5", Dark: "#818CF8"}). + Padding(0, 1) + + TableCellStyle = lipgloss.NewStyle(). + Foreground(lipgloss.AdaptiveColor{Light: "#1E293B", Dark: "#F1F5F9"}). + Padding(0, 1) + + TableNilStyle = lipgloss.NewStyle(). + Foreground(lipgloss.AdaptiveColor{Light: "#94A3B8", Dark: "#64748B"}). + Italic(true). + Padding(0, 1) + + TableBorderStyle = lipgloss.NewStyle(). + Foreground(lipgloss.AdaptiveColor{Light: "#CBD5E1", Dark: "#334155"}) + + ActiveTableBorderStyle = lipgloss.NewStyle(). + Foreground(lipgloss.AdaptiveColor{Light: "#6366F1", Dark: "#818CF8"}) + + FieldKeyStyle = lipgloss.NewStyle(). + Bold(true). + Foreground(lipgloss.AdaptiveColor{Light: "#4F46E5", Dark: "#818CF8"}) + + FieldValueStyle = lipgloss.NewStyle(). + Foreground(lipgloss.AdaptiveColor{Light: "#1E293B", Dark: "#F1F5F9"}) + + RecordDividerStyle = lipgloss.NewStyle(). + Bold(true). + Foreground(PrimaryColor) + + ScrollBadgeStyle = lipgloss.NewStyle(). + Bold(true). + Foreground(lipgloss.AdaptiveColor{Light: "#BE185D", Dark: "#F472B6"}) +) + +const ( + MaxColumnWidth = 24 // Max character width per cell to keep grid compact + PageRowSize = 12 // Rows per page +) + +// FormatTableResult renders a beautiful terminal box table for SQL query results with column scrolling and row pagination. +func FormatTableResult(result *db.QueryResult, colOffset int, rowOffset int, termWidth int, isActive bool) string { + if result == nil || len(result.Columns) == 0 { + return lipgloss.NewStyle().Foreground(MutedColor).Italic(true).Render("(No columns or empty dataset returned)") + } + + if len(result.Rows) == 0 { + return lipgloss.NewStyle().Foreground(MutedColor).Italic(true).Render("(0 rows returned)") + } + + if termWidth <= 20 { + termWidth = 80 + } + + totalRows := len(result.Rows) + if rowOffset >= totalRows { + rowOffset = (totalRows - 1) / PageRowSize * PageRowSize + } + if rowOffset < 0 { + rowOffset = 0 + } + + // Calculate maximum display width needed for each column using runewidth for CJK support + totalCols := len(result.Columns) + if colOffset >= totalCols { + colOffset = totalCols - 1 + } + if colOffset < 0 { + colOffset = 0 + } + + colWidths := make([]int, totalCols) + for i, col := range result.Columns { + w := runewidth.StringWidth(col) + if w > MaxColumnWidth { + w = MaxColumnWidth + } + // Inspect current page rows for width calculation + endR := rowOffset + PageRowSize + if endR > totalRows { + endR = totalRows + } + for r := rowOffset; r < endR; r++ { + val := result.Rows[r][col] + if val != nil { + cellStr := fmt.Sprintf("%v", val) + cellStr = strings.ReplaceAll(cellStr, "\n", " ") + dispLen := runewidth.StringWidth(cellStr) + if dispLen > w { + w = dispLen + } + } + } + if w > MaxColumnWidth { + w = MaxColumnWidth + } + if w < 6 { + w = 6 + } + // Add padding (2 chars) + border (1 char) + colWidths[i] = w + 3 + } + + // Determine visible column range [startCol, endCol) that fits within termWidth - 8 + availWidth := termWidth - 8 + if availWidth < 30 { + availWidth = 30 + } + + startCol := colOffset + endCol := startCol + accumWidth := 0 + + for i := startCol; i < totalCols; i++ { + if accumWidth+colWidths[i] > availWidth && endCol > startCol { + break + } + accumWidth += colWidths[i] + endCol = i + 1 + } + + visibleCols := result.Columns[startCol:endCol] + + // Truncate and clean visible headers + headers := make([]string, len(visibleCols)) + for i, col := range visibleCols { + colIdx := startCol + i + headerText := sanitizeCell(col, colWidths[colIdx]-3) + if i == 0 && startCol > 0 { + headerText = "◀ " + headerText + } + if i == len(visibleCols)-1 && endCol < totalCols { + headerText = headerText + fmt.Sprintf(" ▶(+%d)", totalCols-endCol) + } + headers[i] = headerText + } + + borderStyle := TableBorderStyle + if isActive { + borderStyle = ActiveTableBorderStyle + } + + t := table.New(). + Border(lipgloss.RoundedBorder()). + BorderStyle(borderStyle). + Headers(headers...) + + // Configure header styling + t.StyleFunc(func(row, col int) lipgloss.Style { + if row == table.HeaderRow { + return TableHeaderStyle + } + return TableCellStyle + }) + + // Add page rows [rowOffset, min(totalRows, rowOffset+PageRowSize)) + endRow := rowOffset + PageRowSize + if endRow > totalRows { + endRow = totalRows + } + + hasTruncatedCell := false + + for i := rowOffset; i < endRow; i++ { + var rowValues []string + for idx, col := range visibleCols { + colIdx := startCol + idx + val := result.Rows[i][col] + if val == nil { + rowValues = append(rowValues, TableNilStyle.Render("NULL")) + } else { + cellStr, wasTruncated := sanitizeCellWithStatus(val, colWidths[colIdx]-3) + if wasTruncated { + hasTruncatedCell = true + } + rowValues = append(rowValues, cellStr) + } + } + t.Row(rowValues...) + } + + var sb strings.Builder + sb.WriteString(t.Render()) + + var footerNotes []string + if isActive { + footerNotes = append(footerNotes, "[FOCUSED]") + } + if totalRows > PageRowSize { + footerNotes = append(footerNotes, fmt.Sprintf("rows %d-%d of %d (Press PgUp/PgDn for Page)", rowOffset+1, endRow, totalRows)) + } + if startCol > 0 || endCol < totalCols { + footerNotes = append(footerNotes, fmt.Sprintf("cols %d-%d of %d (Use ←/→ keys for Cols)", startCol+1, endCol, totalCols)) + } + if hasTruncatedCell { + footerNotes = append(footerNotes, "press Ctrl+E for Full View (Expand/Collapse)") + } + + if len(footerNotes) > 0 { + noteStr := "\n(" + strings.Join(footerNotes, " | ") + ")" + if isActive { + sb.WriteString(lipgloss.NewStyle().Foreground(lipgloss.Color("#7AA2F7")).Bold(true).Render(noteStr)) + } else { + sb.WriteString(lipgloss.NewStyle().Foreground(MutedColor).Italic(true).Render(noteStr)) + } + } + + return sb.String() +} + +// FormatVerticalResult renders SQL query results in full vertical (psql \x expanded) format with NO truncation. +func FormatVerticalResult(result *db.QueryResult) string { + if result == nil || len(result.Columns) == 0 { + return lipgloss.NewStyle().Foreground(MutedColor).Italic(true).Render("(No columns or empty dataset returned)") + } + + if len(result.Rows) == 0 { + return lipgloss.NewStyle().Foreground(MutedColor).Italic(true).Render("(0 rows returned)") + } + + maxKeyLen := 0 + for _, col := range result.Columns { + w := runewidth.StringWidth(col) + if w > maxKeyLen { + maxKeyLen = w + } + } + + maxRows := len(result.Rows) + if maxRows > 500 { + maxRows = 500 + } + + var sb strings.Builder + for i := 0; i < maxRows; i++ { + divider := fmt.Sprintf("━ Record %d of %d ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━", i+1, len(result.Rows)) + sb.WriteString(RecordDividerStyle.Render(divider) + "\n") + + for _, col := range result.Columns { + val := result.Rows[i][col] + keyWidth := runewidth.StringWidth(col) + padding := strings.Repeat(" ", max(0, maxKeyLen-keyWidth)) + keyStr := FieldKeyStyle.Render(col + padding) + + if val == nil { + fmt.Fprintf(&sb, " %s : %s\n", keyStr, TableNilStyle.Render("NULL")) + } else { + valStr := fmt.Sprintf("%v", val) + // Full display with indentation for multiline text + if strings.Contains(valStr, "\n") { + indented := strings.ReplaceAll(valStr, "\n", "\n ") + fmt.Fprintf(&sb, " %s :\n %s\n", keyStr, FieldValueStyle.Render(indented)) + } else { + fmt.Fprintf(&sb, " %s : %s\n", keyStr, FieldValueStyle.Render(valStr)) + } + } + } + sb.WriteString("\n") + } + + if len(result.Rows) > 500 { + sb.WriteString(lipgloss.NewStyle().Foreground(MutedColor).Italic(true).Render(fmt.Sprintf("... and %d more rows (truncated for performance)\n", len(result.Rows)-500))) + } + + return sb.String() +} + +func sanitizeCell(val any, maxLen int) string { + res, _ := sanitizeCellWithStatus(val, maxLen) + return res +} + +func sanitizeCellWithStatus(val any, maxLen int) (string, bool) { + if val == nil { + return "NULL", false + } + s := fmt.Sprintf("%v", val) + // Replace all line breaks with spaces so the box border line never breaks + s = strings.ReplaceAll(s, "\r\n", " ") + s = strings.ReplaceAll(s, "\n", " ") + s = strings.ReplaceAll(s, "\r", " ") + s = strings.TrimSpace(s) + + dispLen := runewidth.StringWidth(s) + if dispLen > maxLen { + if maxLen <= 3 { + return runewidth.Truncate(s, maxLen, ""), true + } + return runewidth.Truncate(s, maxLen, "..."), true + } + return s, false +} diff --git a/internal/tui/components_test.go b/internal/tui/components_test.go new file mode 100644 index 0000000..b2be21d --- /dev/null +++ b/internal/tui/components_test.go @@ -0,0 +1,77 @@ +package tui + +import ( + "strings" + "testing" + + "github.com/mattn/go-runewidth" + "github.com/zx06/xsql/internal/db" +) + +func TestSanitizeCellWithCJK(t *testing.T) { + // Test Chinese wide character display width calculation & truncation + cjkStr := "机器下架" + if runewidth.StringWidth(cjkStr) != 8 { + t.Fatalf("expected runewidth 8 for '机器下架', got %d", runewidth.StringWidth(cjkStr)) + } + + sanitized, wasTruncated := sanitizeCellWithStatus(cjkStr, 6) + if !wasTruncated { + t.Fatal("expected wasTruncated to be true") + } + if strings.Contains(sanitized, "\n") { + t.Fatal("sanitized string must never contain newlines") + } +} + +func TestFormatTableResult_CJKBorderProtection(t *testing.T) { + res := &db.QueryResult{ + Columns: []string{"id", "status"}, + Rows: []map[string]any{ + {"id": "1", "status": "故障"}, + {"id": "2", "status": "机器下架"}, + {"id": "3", "status": "正常运行"}, + }, + } + + formatted := FormatTableResult(res, 0, 0, 80, true) + lines := strings.Split(formatted, "\n") + + // Verify that rows are strictly single line per data row + for i, line := range lines { + if strings.Contains(line, "故障") { + if !strings.Contains(line, "1") { + t.Errorf("line %d split '故障' into new line: %q", i, line) + } + } + } + + // Nil and empty result fallbacks + if nilOut := FormatTableResult(nil, 0, 0, 80, false); !strings.Contains(nilOut, "empty dataset") { + t.Errorf("expected empty dataset message for nil QueryResult, got %q", nilOut) + } + + if nilVert := FormatVerticalResult(nil); !strings.Contains(nilVert, "empty dataset") { + t.Errorf("expected empty dataset message for nil QueryResult vertical view, got %q", nilVert) + } + + emptyRes := &db.QueryResult{Columns: []string{"id"}, Rows: []map[string]any{}} + if emptyOut := FormatTableResult(emptyRes, 0, 0, 80, false); !strings.Contains(emptyOut, "0 rows returned") { + t.Fatalf("expected '0 rows returned', got %q", emptyOut) + } + + // Test offset and inactive view + offsetFormatted := FormatTableResult(res, 1, 1, 40, false) + if !strings.Contains(offsetFormatted, "status") { + t.Errorf("expected status column in offset view, got:\n%s", offsetFormatted) + } + + // Test sanitizeCell & sanitizeCellWithStatus + if s, _ := sanitizeCellWithStatus(nil, 10); s != "NULL" { + t.Errorf("expected 'NULL' for nil input, got %q", s) + } + + if s := sanitizeCell("short", 10); s != "short" { + t.Errorf("expected 'short', got %q", s) + } +} diff --git a/internal/tui/model.go b/internal/tui/model.go new file mode 100644 index 0000000..3ca166a --- /dev/null +++ b/internal/tui/model.go @@ -0,0 +1,1061 @@ +package tui + +import ( + "context" + "fmt" + "sort" + "strings" + "time" + + "github.com/charmbracelet/bubbles/spinner" + "github.com/charmbracelet/bubbles/textarea" + "github.com/charmbracelet/bubbles/viewport" + tea "github.com/charmbracelet/bubbletea" + "github.com/charmbracelet/lipgloss" + + "github.com/zx06/xsql/internal/ai" + "github.com/zx06/xsql/internal/app" + "github.com/zx06/xsql/internal/config" + "github.com/zx06/xsql/internal/db" + "github.com/zx06/xsql/internal/errors" + "github.com/zx06/xsql/internal/export" + "github.com/zx06/xsql/internal/js" + "github.com/zx06/xsql/internal/session" +) + +type State int + +const ( + StateLoadingSchema State = iota + StateIdle + StateThinking + StateSQLReady + StateExecuting + StateExportReady +) + +// Msg types +type schemaLoadedMsg struct { + schema *db.SchemaInfo + err *errors.XError +} + +type aiResponseMsg struct { + response *ai.AIResponse + err *errors.XError +} + +type queryExecutedMsg struct { + result *db.QueryResult + err *errors.XError + duration time.Duration +} + +type TableState struct { + Result *db.QueryResult + MsgIndex int + ColOffset int + RowOffset int + VerticalView bool +} + +type ToolCallItem struct { + ID string + Name string + Summary string + Detail string + Result string + RawOutput string // Raw execution output/logs (never hidden!) + TableStateIndex int // -1 if no table attached + MsgIndex int + IsExpanded bool +} + +type PendingExport struct { + DatasetID string + Format string + FilePath string + ToolIdx int +} + +type Model struct { + opts config.Options + aiService *ai.Service + profile config.Profile + profileName string + allProfiles map[string]config.Profile + profileList []string + unsafeAllowWrite bool + initialPrompt string + autoExecute bool + + sessionStore *session.SessionDataStore + jsEngine *js.JSEngine + chatHistory []ai.ChatMessage + pendingExport *PendingExport + jsRetryCount int + maxJSRetries int + lastCtrlCTime time.Time + + confirmOption int // 0: Confirm/Execute, 1: Adjust Prompt, 2: Cancel/Deny + + state State + schemaInfo *db.SchemaInfo + currentSQL string + explanation string + messages []string + tableStates []TableState + toolCalls []ToolCallItem + activeTable int + activeToolIdx int + + textarea textarea.Model + viewport viewport.Model + spinner spinner.Model + + width int + height int +} + +func NewModel(opts config.Options, resolved config.Resolved, aiService *ai.Service, initialPrompt string, unsafeAllowWrite bool) Model { + ta := textarea.New() + ta.Placeholder = "Ask AI to query database or perform data analysis..." + ta.ShowLineNumbers = false + ta.Prompt = "" + ta.Focus() + ta.CharLimit = 4000 + ta.SetWidth(80) + ta.SetHeight(3) + + // Custom crisp styles for textarea + ta.FocusedStyle.CursorLine = lipgloss.NewStyle() + + vp := viewport.New(80, 15) + + s := spinner.New() + s.Spinner = spinner.Dot + s.Style = lipgloss.NewStyle().Foreground(PrimaryColor) + + var pList []string + if len(resolved.AllProfiles) > 0 { + for name := range resolved.AllProfiles { + pList = append(pList, name) + } + sort.Strings(pList) + } else if resolved.ProfileName != "" { + pList = []string{resolved.ProfileName} + } + + return Model{ + opts: opts, + aiService: aiService, + profile: resolved.Profile, + profileName: resolved.ProfileName, + allProfiles: resolved.AllProfiles, + profileList: pList, + unsafeAllowWrite: unsafeAllowWrite || resolved.Profile.UnsafeAllowWrite, + initialPrompt: strings.TrimSpace(initialPrompt), + autoExecute: false, + sessionStore: session.NewSessionDataStore(), + jsEngine: js.NewJSEngine(1 * time.Minute), + chatHistory: []ai.ChatMessage{}, + jsRetryCount: 0, + maxJSRetries: 3, + confirmOption: 0, + tableStates: []TableState{}, + toolCalls: []ToolCallItem{}, + activeTable: -1, + activeToolIdx: -1, + state: StateLoadingSchema, + textarea: ta, + viewport: vp, + spinner: s, + messages: []string{}, + width: 80, + height: 24, + } +} + +func (m Model) Init() tea.Cmd { + return tea.Batch( + m.spinner.Tick, + m.loadSchemaCmd(), + ) +} + +func (m Model) loadSchemaCmd() tea.Cmd { + return func() tea.Msg { + ctx := context.Background() + info, xe := app.DumpSchema(ctx, app.SchemaDumpRequest{ + Profile: m.profile, + AllowPlaintext: m.profile.AllowPlaintext, + SkipHostKeyCheck: m.profile.SSHConfig != nil && m.profile.SSHConfig.SkipHostKey, + }) + return schemaLoadedMsg{schema: info, err: xe} + } +} + +func (m Model) runAgentStepCmd() tea.Cmd { + return func() tea.Msg { + ctx := context.Background() + catalog := m.sessionStore.GetCatalog() + sysPrompt := ai.BuildSystemPrompt(m.profile.DB, m.schemaInfo, catalog) + + msgs := make([]ai.ChatMessage, 0, len(m.chatHistory)+1) + msgs = append(msgs, ai.ChatMessage{Role: "system", Content: sysPrompt}) + for _, item := range m.chatHistory { + if item.Role != "system" { + msgs = append(msgs, item) + } + } + + resp, xe := m.aiService.ChatCompletion(ctx, msgs) + return aiResponseMsg{response: resp, err: xe} + } +} + +func (m Model) executeSQLCmd(sqlStr string) tea.Cmd { + return func() tea.Msg { + start := time.Now() + ctx := context.Background() + res, xe := app.Query(ctx, app.QueryRequest{ + Profile: m.profile, + SQL: sqlStr, + AllowPlaintext: m.profile.AllowPlaintext, + SkipHostKeyCheck: m.profile.SSHConfig != nil && m.profile.SSHConfig.SkipHostKey, + UnsafeAllowWrite: m.unsafeAllowWrite, + }) + elapsed := time.Since(start) + return queryExecutedMsg{result: res, err: xe, duration: elapsed} + } +} + +func (m *Model) focusToolCall(idx int) { + if len(m.toolCalls) == 0 { + return + } + if idx < 0 { + idx = len(m.toolCalls) - 1 + } else if idx >= len(m.toolCalls) { + idx = 0 + } + + oldIdx := m.activeToolIdx + m.activeToolIdx = idx + + // Sync embedded table focus if tool call has attached table + if tc := m.toolCalls[idx]; tc.TableStateIndex >= 0 { + m.activeTable = tc.TableStateIndex + } + + if oldIdx >= 0 && oldIdx < len(m.toolCalls) { + m.renderToolCall(oldIdx) + } + m.renderToolCall(m.activeToolIdx) +} + +func (m *Model) renderTableState(idx int, isActive bool) { + if idx < 0 || idx >= len(m.tableStates) { + return + } + ts := &m.tableStates[idx] + + // Find associated ToolCallItem if embedded + for i := range m.toolCalls { + if m.toolCalls[i].TableStateIndex == idx { + m.renderToolCall(i) + return + } + } + + if ts.MsgIndex < 0 || ts.MsgIndex >= len(m.messages) { + return + } + formatted := FormatTableResult(ts.Result, ts.ColOffset, ts.RowOffset, m.width, isActive) + if ts.VerticalView { + formatted = FormatVerticalResult(ts.Result) + } + m.messages[ts.MsgIndex] = formatted + m.viewport.SetContent(strings.Join(m.messages, "\n\n")) +} + +func (m *Model) renderToolCall(idx int) { + if idx < 0 || idx >= len(m.toolCalls) { + return + } + tc := &m.toolCalls[idx] + if tc.MsgIndex < 0 || tc.MsgIndex >= len(m.messages) { + return + } + + isActiveTool := (idx == m.activeToolIdx) + activeMarker := "" + if isActiveTool && len(m.toolCalls) > 1 { + activeMarker = fmt.Sprintf(" 🎯[Tool %d/%d]", idx+1, len(m.toolCalls)) + } + + var sb strings.Builder + if !tc.IsExpanded { + badge := ToolCollapsedBadge.Render("▶ 🛠️ Tool: " + tc.Name + activeMarker) + summary := MetricsStyle.Render(fmt.Sprintf("%s (Folded - Press Ctrl+O to unfold)", tc.Summary)) + fmt.Fprintf(&sb, "%s %s", badge, summary) + } else { + badge := ToolExpandedBadge.Render("▼ 🛠️ Tool: " + tc.Name + activeMarker) + + detailCode := tc.Detail + switch tc.Name { + case "execute_sql": + detailCode = HighlightSQL(tc.Detail) + case "execute_javascript": + detailCode = HighlightJS(tc.Detail) + } + + detail := ToolDetailStyle.Render(detailCode) + resText := MetricsStyle.Render(tc.Result) + fmt.Fprintf(&sb, "%s\n%s\n%s", badge, detail, resText) + + // Render Raw Output / Calculation Results if present (Never hide tool output!) + if tc.RawOutput != "" { + outTitle := lipgloss.NewStyle().Bold(true).Foreground(SecondaryColor).Render("📊 Raw Execution Output:") + outBox := ToolDetailStyle.Render(tc.RawOutput) + fmt.Fprintf(&sb, "\n%s\n%s", outTitle, outBox) + } + + // Render embedded Table Result inside container when unfolded + if tc.TableStateIndex >= 0 && tc.TableStateIndex < len(m.tableStates) { + ts := &m.tableStates[tc.TableStateIndex] + tableStr := FormatTableResult(ts.Result, ts.ColOffset, ts.RowOffset, m.width, tc.TableStateIndex == m.activeTable) + if ts.VerticalView { + tableStr = FormatVerticalResult(ts.Result) + } + sb.WriteString("\n\n" + tableStr) + } + } + + m.messages[tc.MsgIndex] = sb.String() + m.viewport.SetContent(strings.Join(m.messages, "\n\n")) +} + +func (m Model) Update(msg tea.Msg) (tea.Model, tea.Cmd) { + var cmds []tea.Cmd + + switch msg := msg.(type) { + case tea.WindowSizeMsg: + m.width = msg.Width + m.height = msg.Height + m.textarea.SetWidth(max(20, msg.Width-6)) + m.viewport.Width = max(20, msg.Width-4) + m.viewport.Height = max(5, msg.Height-15) + + case schemaLoadedMsg: + if msg.err != nil { + m.messages = append(m.messages, ErrorMsgStyle.Render(fmt.Sprintf("Failed to load schema for profile '%s': %v", m.profileName, msg.err))) + } else { + m.schemaInfo = msg.schema + } + if m.initialPrompt != "" { + prompt := m.initialPrompt + m.initialPrompt = "" + + userHeader := UserTagStyle.Render("👤 YOU") + var userLine string + if strings.Contains(prompt, "\n") { + userLine = fmt.Sprintf("%s\n%s", userHeader, prompt) + } else { + userLine = fmt.Sprintf("%s %s", userHeader, prompt) + } + + m.messages = append(m.messages, userLine) + m.chatHistory = append(m.chatHistory, ai.ChatMessage{Role: "user", Content: prompt}) + m.state = StateThinking + m.viewport.SetContent(strings.Join(m.messages, "\n\n")) + m.viewport.GotoBottom() + return m, m.runAgentStepCmd() + } + m.state = StateIdle + m.viewport.SetContent(strings.Join(m.messages, "\n\n")) + + case aiResponseMsg: + if msg.err != nil { + m.messages = append(m.messages, ErrorMsgStyle.Render(fmt.Sprintf("AI Error: %v", msg.err))) + m.state = StateIdle + } else { + m.explanation = msg.response.Explanation + + if msg.response.Type == ai.TypeJS && msg.response.JSCode != "" { + m.chatHistory = append(m.chatHistory, ai.ChatMessage{ + Role: "assistant", + Content: fmt.Sprintf("Call tool 'execute_javascript':\n%s", msg.response.JSCode), + }) + + lineCount := len(strings.Split(msg.response.JSCode, "\n")) + tc := ToolCallItem{ + ID: fmt.Sprintf("tc_%d", len(m.toolCalls)+1), + Name: "execute_javascript", + Summary: fmt.Sprintf("Executing %d lines of JS data analysis", lineCount), + Detail: msg.response.JSCode, + TableStateIndex: -1, + MsgIndex: len(m.messages), + IsExpanded: false, + } + + m.messages = append(m.messages, "") + m.toolCalls = append(m.toolCalls, tc) + toolIdx := len(m.toolCalls) - 1 + m.focusToolCall(toolIdx) + + ctx := context.Background() + jsRes, jsErr := m.jsEngine.Execute(ctx, msg.response.JSCode, m.sessionStore) + if jsErr != nil { + m.jsRetryCount++ + m.toolCalls[toolIdx].Result = fmt.Sprintf("❌ Failed: %v", jsErr.Message) + m.renderToolCall(toolIdx) + + if m.jsRetryCount <= m.maxJSRetries { + retryWarn := ErrorMsgStyle.Render(fmt.Sprintf("⚠️ JS Execution Failed (Attempt %d/%d): %v", m.jsRetryCount, m.maxJSRetries, jsErr.Message)) + m.messages = append(m.messages, retryWarn) + + m.chatHistory = append(m.chatHistory, ai.ChatMessage{ + Role: "user", + Content: fmt.Sprintf("Tool 'execute_javascript' failed with error: %s. Please fix the code and call 'execute_javascript' again.", jsErr.Message), + }) + + m.state = StateThinking + m.viewport.SetContent(strings.Join(m.messages, "\n\n")) + m.viewport.GotoBottom() + return m, m.runAgentStepCmd() + } + m.messages = append(m.messages, ErrorMsgStyle.Render(fmt.Sprintf("❌ JS Execution Error (after %d retries): %v", m.maxJSRetries, jsErr.Message))) + m.jsRetryCount = 0 + m.state = StateIdle + } else { + m.jsRetryCount = 0 + m.toolCalls[toolIdx].Result = "✓ JavaScript executed successfully" + m.toolCalls[toolIdx].RawOutput = jsRes.SummaryText // Display raw JS output inside tool call container! + m.renderToolCall(toolIdx) + + m.chatHistory = append(m.chatHistory, ai.ChatMessage{ + Role: "user", + Content: fmt.Sprintf("Tool 'execute_javascript' executed successfully. Output:\n%s", jsRes.SummaryText), + }) + + m.state = StateThinking + m.viewport.SetContent(strings.Join(m.messages, "\n\n")) + m.viewport.GotoBottom() + return m, m.runAgentStepCmd() + } + } else if msg.response.Type == ai.TypeExport && msg.response.DatasetID != "" { + m.chatHistory = append(m.chatHistory, ai.ChatMessage{ + Role: "assistant", + Content: fmt.Sprintf("Call tool 'export_data': dataset_id=%s, format=%s, filepath=%s", msg.response.DatasetID, msg.response.Format, msg.response.FilePath), + }) + + tc := ToolCallItem{ + ID: fmt.Sprintf("tc_%d", len(m.toolCalls)+1), + Name: "export_data", + Summary: fmt.Sprintf("Export %s to %s (%s) [Pending User Confirmation]", msg.response.DatasetID, msg.response.FilePath, strings.ToUpper(msg.response.Format)), + Detail: fmt.Sprintf("Dataset: %s | FilePath: %s | Format: %s", msg.response.DatasetID, msg.response.FilePath, msg.response.Format), + Result: "⏳ Pending Human Confirmation", + TableStateIndex: -1, + MsgIndex: len(m.messages), + IsExpanded: false, + } + m.messages = append(m.messages, "") + m.toolCalls = append(m.toolCalls, tc) + toolIdx := len(m.toolCalls) - 1 + m.focusToolCall(toolIdx) + + m.pendingExport = &PendingExport{ + DatasetID: msg.response.DatasetID, + Format: msg.response.Format, + FilePath: msg.response.FilePath, + ToolIdx: toolIdx, + } + m.confirmOption = 0 + m.state = StateExportReady + } else if msg.response.Type == ai.TypeSQL && msg.response.SQL != "" { + m.chatHistory = append(m.chatHistory, ai.ChatMessage{ + Role: "assistant", + Content: fmt.Sprintf("Call tool 'execute_sql': %s", msg.response.SQL), + }) + + m.currentSQL = msg.response.SQL + + tc := ToolCallItem{ + ID: fmt.Sprintf("tc_%d", len(m.toolCalls)+1), + Name: "execute_sql", + Summary: m.currentSQL, + Detail: m.currentSQL, + TableStateIndex: -1, + MsgIndex: len(m.messages), + IsExpanded: false, + } + m.messages = append(m.messages, "") + m.toolCalls = append(m.toolCalls, tc) + toolIdx := len(m.toolCalls) - 1 + m.focusToolCall(toolIdx) + + if m.autoExecute { + m.state = StateExecuting + m.viewport.SetContent(strings.Join(m.messages, "\n\n")) + m.viewport.GotoBottom() + return m, m.executeSQLCmd(m.currentSQL) + } + m.confirmOption = 0 + m.state = StateSQLReady + } else { + // FINAL LLM AGENT OUTPUT (No Tool Call) + exp := msg.response.Explanation + + // Defensive guard: If LLM mistakenly output raw JS code in text instead of tool call, intercept it! + if strings.Contains(exp, "Call tool 'execute_javascript'") || (strings.Contains(exp, "var data =") && strings.Contains(exp, "stats")) { + jsCode := exp + if idx := strings.Index(exp, "var "); idx >= 0 { + jsCode = exp[idx:] + } + msg.response.Type = ai.TypeJS + msg.response.JSCode = jsCode + return m.Update(msg) + } + + m.chatHistory = append(m.chatHistory, ai.ChatMessage{ + Role: "assistant", + Content: exp, + }) + + if exp != "" { + renderedMD := RenderMarkdown(exp, m.width) + aiMsg := AITagStyle.Render("🤖 AI") + "\n" + renderedMD + m.messages = append(m.messages, aiMsg) + } + m.state = StateIdle + } + } + m.viewport.SetContent(strings.Join(m.messages, "\n\n")) + m.viewport.GotoBottom() + + case queryExecutedMsg: + if msg.err != nil { + errText := fmt.Sprintf("SQL Exec Error [%s]: %s", msg.err.Code, msg.err.Message) + m.messages = append(m.messages, ErrorMsgStyle.Render(errText)) + + m.chatHistory = append(m.chatHistory, ai.ChatMessage{ + Role: "user", + Content: fmt.Sprintf("Tool 'execute_sql' failed with error: %s", errText), + }) + + m.state = StateThinking + m.viewport.SetContent(strings.Join(m.messages, "\n\n")) + m.viewport.GotoBottom() + return m, m.runAgentStepCmd() + } else if msg.result != nil { + datasetID := m.sessionStore.Save(m.currentSQL, msg.result) + + modelName := m.opts.CLIAIModel + if modelName == "" { + modelName = "gpt-4o" + } + durStr := msg.duration.Round(time.Millisecond).String() + if msg.duration < time.Millisecond { + durStr = fmt.Sprintf("%.2fms", float64(msg.duration.Microseconds())/1000.0) + } + metricsStr := fmt.Sprintf("⏱️ %s | 📊 %d rows | 🤖 %s | 💾 %s", durStr, len(msg.result.Rows), modelName, datasetID) + statusLine := SuccessBadgeStyle.Render("✓ Execution Success") + " " + MetricsStyle.Render(metricsStr) + + ts := TableState{ + Result: msg.result, + MsgIndex: -1, + ColOffset: 0, + RowOffset: 0, + VerticalView: false, + } + m.tableStates = append(m.tableStates, ts) + tableIdx := len(m.tableStates) - 1 + + // Attach TableStateIndex directly inside execute_sql ToolCallItem + if len(m.toolCalls) > 0 { + lastIdx := len(m.toolCalls) - 1 + if m.toolCalls[lastIdx].Name == "execute_sql" { + m.toolCalls[lastIdx].Result = statusLine + m.toolCalls[lastIdx].TableStateIndex = tableIdx + m.focusToolCall(lastIdx) + } + } + + m.chatHistory = append(m.chatHistory, ai.ChatMessage{ + Role: "user", + Content: fmt.Sprintf("Tool 'execute_sql' executed successfully. Returned %d rows (columns: %v). Dataset saved as '%s'.", len(msg.result.Rows), msg.result.Columns, datasetID), + }) + + m.state = StateThinking + m.viewport.SetContent(strings.Join(m.messages, "\n\n")) + m.viewport.GotoBottom() + return m, m.runAgentStepCmd() + } + m.state = StateIdle + m.viewport.SetContent(strings.Join(m.messages, "\n\n")) + m.viewport.GotoBottom() + + case spinner.TickMsg: + var cmd tea.Cmd + m.spinner, cmd = m.spinner.Update(msg) + cmds = append(cmds, cmd) + + case tea.KeyMsg: + // CTRL+C TWICE TO QUIT MECHANISM (Like Claude Code / Aider) + if msg.Type == tea.KeyCtrlC { + if time.Since(m.lastCtrlCTime) < 2*time.Second { + return m, tea.Quit + } + m.lastCtrlCTime = time.Now() + warnMsg := WarningBadgeStyle.Render("⚠️ Press Ctrl+C again to exit xsql AI") + m.messages = append(m.messages, warnMsg) + m.viewport.SetContent(strings.Join(m.messages, "\n\n")) + m.viewport.GotoBottom() + return m, nil + } + + // UNIFIED HUMAN-IN-THE-LOOP INTERACTION FOR BOTH StateExportReady AND StateSQLReady! + if m.state == StateExportReady && m.pendingExport != nil { + switch msg.Type { + case tea.KeyUp: + m.confirmOption = (m.confirmOption - 1 + 3) % 3 + return m, nil + case tea.KeyDown: + m.confirmOption = (m.confirmOption + 1) % 3 + return m, nil + } + + triggerOpt := -1 + if msg.Type == tea.KeyEnter { + triggerOpt = m.confirmOption + } else if msg.String() == "1" { + triggerOpt = 0 + } else if msg.String() == "2" { + triggerOpt = 1 + } else if msg.String() == "3" || msg.Type == tea.KeyEsc { + triggerOpt = 2 + } + + switch triggerOpt { + case 0: + // Option 1: Confirm & Export + datasetRes, exists := m.sessionStore.Get(m.pendingExport.DatasetID) + if !exists || datasetRes == nil { + m.toolCalls[m.pendingExport.ToolIdx].Result = fmt.Sprintf("❌ Export Failed: Dataset '%s' not found", m.pendingExport.DatasetID) + m.renderToolCall(m.pendingExport.ToolIdx) + m.chatHistory = append(m.chatHistory, ai.ChatMessage{ + Role: "user", + Content: fmt.Sprintf("Tool 'export_data' failed: dataset '%s' not found in session catalog.", m.pendingExport.DatasetID), + }) + } else { + outPath, xe := export.ExportQueryResult(datasetRes, export.ExportFormat(m.pendingExport.Format), m.pendingExport.FilePath) + if xe != nil { + m.toolCalls[m.pendingExport.ToolIdx].Result = fmt.Sprintf("❌ Export Failed: %v", xe.Message) + m.renderToolCall(m.pendingExport.ToolIdx) + m.chatHistory = append(m.chatHistory, ai.ChatMessage{ + Role: "user", + Content: fmt.Sprintf("Tool 'export_data' failed to write file: %v", xe.Message), + }) + } else { + m.toolCalls[m.pendingExport.ToolIdx].Result = fmt.Sprintf("✓ Exported dataset '%s' to '%s' (%s)", m.pendingExport.DatasetID, outPath, strings.ToUpper(m.pendingExport.Format)) + m.renderToolCall(m.pendingExport.ToolIdx) + + statusLine := SuccessBadgeStyle.Render("✓ File Exported Success") + " " + MetricsStyle.Render(fmt.Sprintf("Exported dataset '%s' to '%s'", m.pendingExport.DatasetID, outPath)) + m.messages = append(m.messages, statusLine) + + m.chatHistory = append(m.chatHistory, ai.ChatMessage{ + Role: "user", + Content: fmt.Sprintf("Tool 'export_data' executed successfully. Exported dataset '%s' to local file '%s'.", m.pendingExport.DatasetID, outPath), + }) + } + } + m.pendingExport = nil + m.state = StateThinking + m.viewport.SetContent(strings.Join(m.messages, "\n\n")) + m.viewport.GotoBottom() + return m, m.runAgentStepCmd() + + case 1: + // Option 2: Adjust Prompt + m.state = StateIdle + m.pendingExport = nil + m.textarea.Focus() + return m, nil + + case 2: + // Option 3: Deny / Cancel Export + m.toolCalls[m.pendingExport.ToolIdx].Result = "🚫 Export Denied by User" + m.renderToolCall(m.pendingExport.ToolIdx) + + m.chatHistory = append(m.chatHistory, ai.ChatMessage{ + Role: "user", + Content: "Tool 'export_data' was denied by user.", + }) + m.pendingExport = nil + m.state = StateThinking + m.viewport.SetContent(strings.Join(m.messages, "\n\n")) + m.viewport.GotoBottom() + return m, m.runAgentStepCmd() + } + return m, nil + } + + if m.state == StateSQLReady { + switch msg.Type { + case tea.KeyUp: + m.confirmOption = (m.confirmOption - 1 + 3) % 3 + return m, nil + case tea.KeyDown: + m.confirmOption = (m.confirmOption + 1) % 3 + return m, nil + } + + triggerOpt := -1 + if msg.Type == tea.KeyEnter { + triggerOpt = m.confirmOption + } else if msg.String() == "1" { + triggerOpt = 0 + } else if msg.String() == "2" { + triggerOpt = 1 + } else if msg.String() == "3" || msg.Type == tea.KeyEsc { + triggerOpt = 2 + } + + switch triggerOpt { + case 0: + // Option 1: Execute SQL + m.state = StateExecuting + m.textarea.Focus() + m.viewport.SetContent(strings.Join(m.messages, "\n\n")) + m.viewport.GotoBottom() + return m, m.executeSQLCmd(m.currentSQL) + + case 1: + // Option 2: Adjust Prompt / Re-generate + m.state = StateIdle + m.textarea.Focus() + return m, nil + + case 2: + // Option 3: Cancel Execution + m.state = StateIdle + m.textarea.Focus() + return m, nil + } + } + + // MULTI-LINE PROMPT INPUT SHORTCUT: Alt+Enter to insert soft newline into textarea + if m.state == StateIdle && ((msg.Alt && msg.Type == tea.KeyEnter) || msg.String() == "alt+enter") { + m.textarea.InsertString("\n") + return m, nil + } + + switch msg.Type { + case tea.KeyEsc: + // ESC CLEARS PROMPT INPUT BOX (No longer quits application) + if m.state == StateIdle { + m.textarea.Reset() + return m, nil + } + + case tea.KeyCtrlP: + // PROFILE SWITCHING FEATURE: Cycle active profile & synchronize with AI agent context! + if len(m.profileList) > 1 { + currIdx := -1 + for i, name := range m.profileList { + if name == m.profileName { + currIdx = i + break + } + } + nextIdx := (currIdx + 1) % len(m.profileList) + nextProfileName := m.profileList[nextIdx] + + newP, ok := m.allProfiles[nextProfileName] + if ok { + if newP.Port == 0 { + switch newP.DB { + case "mysql": + newP.Port = 3306 + case "pg": + newP.Port = 5432 + } + } + + m.profileName = nextProfileName + m.profile = newP + m.unsafeAllowWrite = newP.UnsafeAllowWrite + m.schemaInfo = nil + m.state = StateLoadingSchema + + switchLine := SuccessBadgeStyle.Render(fmt.Sprintf("✓ Switched active profile to '%s' (%s)", m.profileName, m.profile.DB)) + m.messages = append(m.messages, switchLine) + + m.chatHistory = append(m.chatHistory, ai.ChatMessage{ + Role: "user", + Content: fmt.Sprintf("System Notice: Switched active database profile to '%s' (Database: %s). New database schema metadata is being loaded.", m.profileName, m.profile.DB), + }) + + m.viewport.SetContent(strings.Join(m.messages, "\n\n")) + m.viewport.GotoBottom() + return m, m.loadSchemaCmd() + } + } + + case tea.KeyCtrlO: + // Toggle folding/unfolding of currently active/focused ToolCallItem + if len(m.toolCalls) > 0 { + if m.activeToolIdx < 0 || m.activeToolIdx >= len(m.toolCalls) { + m.activeToolIdx = len(m.toolCalls) - 1 + } + m.toolCalls[m.activeToolIdx].IsExpanded = !m.toolCalls[m.activeToolIdx].IsExpanded + m.renderToolCall(m.activeToolIdx) + return m, nil + } + + case tea.KeyTab: + // Tab: Cycle focus between Tool Call Containers + if len(m.toolCalls) > 0 { + m.focusToolCall((m.activeToolIdx + 1) % len(m.toolCalls)) + return m, nil + } + + case tea.KeyShiftTab: + // Shift+Tab: Toggle AUTO-EXECUTE / MANUAL approval mode + m.autoExecute = !m.autoExecute + + case tea.KeyCtrlE: + if m.activeTable >= 0 && m.activeTable < len(m.tableStates) { + ts := &m.tableStates[m.activeTable] + ts.VerticalView = !ts.VerticalView + m.renderTableState(m.activeTable, true) + } + + case tea.KeyLeft: + if m.activeTable >= 0 && m.activeTable < len(m.tableStates) { + ts := &m.tableStates[m.activeTable] + if ts.ColOffset > 0 { + ts.ColOffset-- + m.renderTableState(m.activeTable, true) + } + } + + case tea.KeyRight: + if m.activeTable >= 0 && m.activeTable < len(m.tableStates) { + ts := &m.tableStates[m.activeTable] + if ts.ColOffset < len(ts.Result.Columns)-1 { + ts.ColOffset++ + m.renderTableState(m.activeTable, true) + } + } + + case tea.KeyPgUp: + if m.activeTable >= 0 && m.activeTable < len(m.tableStates) { + ts := &m.tableStates[m.activeTable] + if ts.RowOffset >= PageRowSize { + ts.RowOffset -= PageRowSize + m.renderTableState(m.activeTable, true) + return m, nil + } + } + m.viewport.LineUp(6) + return m, nil + + case tea.KeyPgDown: + if m.activeTable >= 0 && m.activeTable < len(m.tableStates) { + ts := &m.tableStates[m.activeTable] + if ts.RowOffset+PageRowSize < len(ts.Result.Rows) { + ts.RowOffset += PageRowSize + m.renderTableState(m.activeTable, true) + return m, nil + } + } + m.viewport.LineDown(6) + return m, nil + + case tea.KeyUp: + m.viewport.LineUp(1) + return m, nil + + case tea.KeyDown: + m.viewport.LineDown(1) + return m, nil + + case tea.KeyEnter: + prompt := strings.TrimSpace(m.textarea.Value()) + if prompt != "" && m.state == StateIdle { + userHeader := UserTagStyle.Render("👤 YOU") + var userLine string + if strings.Contains(prompt, "\n") { + userLine = fmt.Sprintf("%s\n%s", userHeader, prompt) + } else { + userLine = fmt.Sprintf("%s %s", userHeader, prompt) + } + + m.messages = append(m.messages, userLine) + m.chatHistory = append(m.chatHistory, ai.ChatMessage{Role: "user", Content: prompt}) + m.textarea.Reset() + m.state = StateThinking + m.viewport.SetContent(strings.Join(m.messages, "\n\n")) + m.viewport.GotoBottom() + return m, m.runAgentStepCmd() + } + } + } + + var taCmd tea.Cmd + m.textarea, taCmd = m.textarea.Update(msg) + cmds = append(cmds, taCmd) + + var vpCmd tea.Cmd + m.viewport, vpCmd = m.viewport.Update(msg) + cmds = append(cmds, vpCmd) + + return m, tea.Batch(cmds...) +} + +func renderKeybindingBadges(items [][2]string) string { + var parts []string + for _, item := range items { + keyBadge := KeyBadgeStyle.Render(item[0]) + label := KeyLabelStyle.Render(item[1]) + parts = append(parts, fmt.Sprintf("%s %s", keyBadge, label)) + } + return strings.Join(parts, " ") +} + +func renderActionOptionsCard(title string, detailText string, options []string, activeOpt int, width int) string { + var sb strings.Builder + sb.WriteString(SQLTitleStyle.Render(title) + "\n") + if detailText != "" { + sb.WriteString(detailText + "\n\n") + } + sb.WriteString(lipgloss.NewStyle().Foreground(MutedColor).Render("Action Options (Use ↑/↓ or 1/2/3 to select, Enter to confirm):") + "\n") + for i, opt := range options { + if i == activeOpt { + sb.WriteString(lipgloss.NewStyle().Bold(true).Foreground(PrimaryColor).Render(fmt.Sprintf(" ▶ [%d] %s", i+1, opt)) + "\n") + } else { + sb.WriteString(lipgloss.NewStyle().Foreground(MutedColor).Render(fmt.Sprintf(" [%d] %s", i+1, opt)) + "\n") + } + } + return SQLBoxStyle.Width(width - 4).Render(sb.String()) +} + +func (m Model) View() string { + var sb strings.Builder + + // 1. Full-Width Header Bar + titlePill := HeaderTitleBadge.Render("xsql AI") + profilePill := HeaderProfileBadge.Render(fmt.Sprintf("%s (%s)", m.profileName, m.profile.DB)) + + modePill := BadgeReadOnly.Render("READ-ONLY") + if m.unsafeAllowWrite { + modePill = BadgeReadWrite.Render("READ-WRITE") + } + + execPill := BadgeManualApprove.Render("MANUAL") + if m.autoExecute { + execPill = BadgeAutoExec.Render("AUTO-EXEC") + } + + headerContent := fmt.Sprintf("%s %s %s %s", titlePill, profilePill, modePill, execPill) + header := HeaderBarStyle.Width(m.width).Render(headerContent) + sb.WriteString(header + "\n\n") + + // 2. Main Viewport + sb.WriteString(m.viewport.View() + "\n\n") + + // 3. State Status & Unified SQL / Export Action Selection Card + switch m.state { + case StateLoadingSchema: + sb.WriteString(m.spinner.View() + fmt.Sprintf(" Loading database schema for profile '%s'...\n", m.profileName)) + case StateThinking: + sb.WriteString(m.spinner.View() + " AI is analyzing schema and executing tools...\n") + case StateExecuting: + sb.WriteString(m.spinner.View() + " Executing SQL query...\n") + case StateExportReady: + if m.pendingExport != nil { + exportInfo := fmt.Sprintf("Dataset: %s | FilePath: %s | Format: %s", m.pendingExport.DatasetID, m.pendingExport.FilePath, strings.ToUpper(m.pendingExport.Format)) + card := renderActionOptionsCard( + "✨ File Export Approval Required", + SQLCodeStyle.Render(exportInfo), + []string{"Confirm & Export File", "Adjust Prompt / Change Options", "Deny & Cancel Export"}, + m.confirmOption, + m.width, + ) + sb.WriteString(card + "\n") + } + case StateSQLReady: + sqlContent := HighlightSQL(m.currentSQL) + if m.currentSQL == "" { + sqlContent = lipgloss.NewStyle().Foreground(MutedColor).Italic(true).Render("(No SQL generated)") + } + card := renderActionOptionsCard( + "✨ SQL Approval Required", + sqlContent, + []string{"Execute SQL Query", "Adjust Prompt / Re-generate", "Cancel Execution"}, + m.confirmOption, + m.width, + ) + sb.WriteString(card + "\n") + } + + // 4. Ultra-Minimal Prompt Input Area & Dynamic Divider Rules + divLen := m.width - 4 + if divLen < 10 { + divLen = 10 + } + separator := lipgloss.NewStyle().Foreground(lipgloss.AdaptiveColor{Light: "#CBD5E1", Dark: "#334155"}).Render(strings.Repeat("─", divLen)) + + sb.WriteString(separator + "\n") + sb.WriteString(m.textarea.View() + "\n") + sb.WriteString(separator + "\n\n") + + execModeHint := "MANUAL" + if m.autoExecute { + execModeHint = "AUTO" + } + + toolFoldState := "Folded" + if m.activeToolIdx >= 0 && m.activeToolIdx < len(m.toolCalls) && m.toolCalls[m.activeToolIdx].IsExpanded { + toolFoldState = "Unfolded" + } + + toolNavHint := "" + if len(m.toolCalls) > 1 { + toolNavHint = fmt.Sprintf(" [%d/%d]", m.activeToolIdx+1, len(m.toolCalls)) + } + + var keybindings string + if m.state == StateExportReady || m.state == StateSQLReady { + keybindings = renderKeybindingBadges([][2]string{ + {"↑/↓", "Select Option"}, + {"Enter", "Confirm"}, + {"1/2/3", "Quick Select"}, + {"Esc", "Cancel"}, + }) + } else { + profileBadge := "" + if len(m.profileList) > 1 { + profileBadge = fmt.Sprintf(" (%s)", m.profileName) + } + keybindings = renderKeybindingBadges([][2]string{ + {"Enter", "Send"}, + {"Alt+Enter", "Newline"}, + {"Ctrl+P", "Profile" + profileBadge}, + {"Tab", "Focus Tool" + toolNavHint}, + {"Ctrl+O", "Tools (" + toolFoldState + ")"}, + {"←/→", "Cols"}, + {"PgUp/PgDn", "Rows"}, + {"Ctrl+E", "Expand Table"}, + {"Shift+Tab", "Mode (" + execModeHint + ")"}, + {"Esc", "Clear"}, + {"Ctrl+C", "Quit (x2)"}, + }) + } + + sb.WriteString(keybindings + "\n") + + return sb.String() +} diff --git a/internal/tui/model_test.go b/internal/tui/model_test.go new file mode 100644 index 0000000..150bcb2 --- /dev/null +++ b/internal/tui/model_test.go @@ -0,0 +1,540 @@ +package tui + +import ( + "strings" + "testing" + "time" + + tea "github.com/charmbracelet/bubbletea" + + "github.com/zx06/xsql/internal/ai" + "github.com/zx06/xsql/internal/config" + "github.com/zx06/xsql/internal/db" + "github.com/zx06/xsql/internal/errors" +) + +func TestTUI_Model_StateTransitions(t *testing.T) { + resolved := config.Resolved{ + ProfileName: "dev", + Profile: config.Profile{ + DB: "mysql", + }, + } + aiService := ai.NewService(config.AIConfig{}, nil) + m := NewModel(config.Options{}, resolved, aiService, "", false) + + // Initial State should be StateLoadingSchema + if m.state != StateLoadingSchema { + t.Fatalf("expected initial state StateLoadingSchema, got %v", m.state) + } + + // 1. Send schemaLoadedMsg -> transition to StateIdle + updated, _ := m.Update(schemaLoadedMsg{ + schema: &db.SchemaInfo{Database: "testdb"}, + }) + m = updated.(Model) + if m.state != StateIdle { + t.Fatalf("expected state StateIdle, got %v", m.state) + } + + // 2. Send aiResponseMsg -> transition to StateSQLReady + updated, _ = m.Update(aiResponseMsg{ + response: &ai.AIResponse{ + Type: ai.TypeSQL, + SQL: "SELECT * FROM users;", + Explanation: "Returns all users.", + }, + }) + m = updated.(Model) + if m.state != StateSQLReady { + t.Fatalf("expected state StateSQLReady, got %v", m.state) + } + if m.currentSQL != "SELECT * FROM users;" { + t.Errorf("expected SQL 'SELECT * FROM users;', got %q", m.currentSQL) + } + + // 3. Test View Output Rendering + viewStr := m.View() + if !strings.Contains(viewStr, "xsql AI") { + t.Errorf("expected view to contain header 'xsql AI', got:\n%s", viewStr) + } + if !strings.Contains(viewStr, "READ-ONLY") { + t.Errorf("expected view to contain READ-ONLY badge, got:\n%s", viewStr) + } + if !strings.Contains(viewStr, "SELECT * FROM users;") { + t.Errorf("expected view to contain SQL Preview, got:\n%s", viewStr) + } + + // 4. Test KeyMsg KeyEnter in StateSQLReady -> transition to StateExecuting + updated, cmd := m.Update(tea.KeyMsg{Type: tea.KeyEnter}) + m = updated.(Model) + if m.state != StateExecuting { + t.Fatalf("expected state StateExecuting after KeyEnter, got %v", m.state) + } + if cmd == nil { + t.Fatal("expected non-nil Cmd for executeSQLCmd") + } + + // 5. Send queryExecutedMsg -> Agent loops back with runAgentStepCmd (StateThinking) + updated, cmd = m.Update(queryExecutedMsg{ + result: &db.QueryResult{ + Columns: []string{"id", "name"}, + Rows: []map[string]any{{"id": 1, "name": "Alice"}}, + }, + }) + m = updated.(Model) + if m.state != StateThinking { + t.Fatalf("expected state StateThinking while Agent processes tool result, got %v", m.state) + } + if cmd == nil { + t.Fatal("expected non-nil Cmd for runAgentStepCmd after query execution") + } + + // 6. Send final aiResponseMsg (TypeText) -> transition to StateIdle + updated, _ = m.Update(aiResponseMsg{ + response: &ai.AIResponse{ + Type: ai.TypeText, + Explanation: "Found 1 user named Alice.", + }, + }) + m = updated.(Model) + if m.state != StateIdle { + t.Fatalf("expected state StateIdle after final text response, got %v", m.state) + } + + // View output should contain query result + viewStr = m.View() + if !strings.Contains(viewStr, "Alice") { + t.Errorf("expected view to contain result 'Alice', got:\n%s", viewStr) + } +} + +func TestFormatTableResult(t *testing.T) { + res := &db.QueryResult{ + Columns: []string{"id", "username", "extra_json"}, + Rows: []map[string]any{ + {"id": 1, "username": "admin", "extra_json": "{\n \"key\": \"very long value that exceeds column limit\"\n}"}, + {"id": 2, "username": "guest", "extra_json": nil}, + }, + } + + formatted := FormatTableResult(res, 0, 0, 80, true) + if !strings.Contains(formatted, "admin") || !strings.Contains(formatted, "guest") { + t.Errorf("formatted table result missing row data:\n%s", formatted) + } + if !strings.Contains(formatted, "NULL") { + t.Errorf("expected NULL representation for nil value, got:\n%s", formatted) + } + if strings.Contains(formatted, "\n \"key\"") { + t.Errorf("expected newlines inside cells to be sanitized, got:\n%s", formatted) + } + + vertFormatted := FormatVerticalResult(res) + if !strings.Contains(vertFormatted, "Record 1 of 2") { + t.Errorf("expected vertical view header, got:\n%s", vertFormatted) + } + if !strings.Contains(vertFormatted, "very long value that exceeds column limit") { + t.Errorf("expected vertical view to display untruncated multiline text, got:\n%s", vertFormatted) + } +} + +func TestTUI_Model_InitialPromptAutoExecute(t *testing.T) { + resolved := config.Resolved{ + ProfileName: "dev", + Profile: config.Profile{DB: "mysql"}, + } + aiService := ai.NewService(config.AIConfig{}, nil) + m := NewModel(config.Options{}, resolved, aiService, "Show total users", false) + + // When schemaLoadedMsg arrives, initialPrompt should automatically trigger StateThinking + updated, cmd := m.Update(schemaLoadedMsg{ + schema: &db.SchemaInfo{Database: "testdb"}, + }) + m = updated.(Model) + + if m.state != StateThinking { + t.Fatalf("expected state StateThinking when initialPrompt provided, got %v", m.state) + } + if cmd == nil { + t.Fatal("expected non-nil Cmd for generateSQLCmd from initial prompt") + } +} + +func TestTUI_Model_ShiftTabAutoExecuteToggle(t *testing.T) { + resolved := config.Resolved{ + ProfileName: "dev", + Profile: config.Profile{DB: "mysql"}, + } + aiService := ai.NewService(config.AIConfig{}, nil) + m := NewModel(config.Options{}, resolved, aiService, "", false) + + if m.autoExecute { + t.Fatal("expected autoExecute to be false by default") + } + + // Press Shift+Tab -> toggle to autoExecute = true + updated, _ := m.Update(tea.KeyMsg{Type: tea.KeyShiftTab}) + m = updated.(Model) + if !m.autoExecute { + t.Fatal("expected autoExecute to be true after Shift+Tab") + } + + // Send aiResponseMsg -> should automatically transition to StateExecuting + updated, cmd := m.Update(aiResponseMsg{ + response: &ai.AIResponse{ + Type: ai.TypeSQL, + SQL: "SELECT * FROM users;", + Explanation: "Returns users.", + }, + }) + m = updated.(Model) + if m.state != StateExecuting { + t.Fatalf("expected state StateExecuting when autoExecute is true, got %v", m.state) + } + if cmd == nil { + t.Fatal("expected non-nil executeSQLCmd for auto-execution") + } +} + +func TestTUI_Model_ActionOptionsCardFlow(t *testing.T) { + resolved := config.Resolved{ + ProfileName: "dev", + Profile: config.Profile{DB: "mysql"}, + } + aiService := ai.NewService(config.AIConfig{}, nil) + m := NewModel(config.Options{}, resolved, aiService, "", false) + m.state = StateSQLReady + m.currentSQL = "SELECT * FROM users LIMIT 10;" + + // 1. Press Down -> switches confirmOption to 1 (Adjust Prompt) + updated, _ := m.Update(tea.KeyMsg{Type: tea.KeyDown}) + m = updated.(Model) + if m.confirmOption != 1 { + t.Fatalf("expected confirmOption to be 1 after KeyDown, got %d", m.confirmOption) + } + + // 2. Press Enter -> Option 1 returns to StateIdle for adjusting prompt + updated, _ = m.Update(tea.KeyMsg{Type: tea.KeyEnter}) + m = updated.(Model) + if m.state != StateIdle { + t.Fatalf("expected state StateIdle after selecting Adjust Prompt option, got %v", m.state) + } + + // 3. Reset to StateSQLReady and press Enter on Option 0 -> transitions to StateExecuting + m.state = StateSQLReady + m.confirmOption = 0 + updated, cmd := m.Update(tea.KeyMsg{Type: tea.KeyEnter}) + m = updated.(Model) + if m.state != StateExecuting { + t.Fatalf("expected state StateExecuting after confirming execution, got %v", m.state) + } + if cmd == nil { + t.Fatal("expected non-nil executeSQLCmd for executing SQL") + } +} + +func TestTUI_Model_CtrlCTwiceToQuit(t *testing.T) { + resolved := config.Resolved{ProfileName: "dev", Profile: config.Profile{DB: "mysql"}} + aiService := ai.NewService(config.AIConfig{}, nil) + m := NewModel(config.Options{}, resolved, aiService, "", false) + + // 1st Ctrl+C -> should NOT quit, but set warning line + updated, cmd := m.Update(tea.KeyMsg{Type: tea.KeyCtrlC}) + m = updated.(Model) + if cmd != nil { + t.Fatal("expected 1st Ctrl+C to NOT return quit cmd") + } + + // 2nd Ctrl+C immediately -> returns tea.Quit + _, cmd = m.Update(tea.KeyMsg{Type: tea.KeyCtrlC}) + if cmd == nil { + t.Fatal("expected 2nd Ctrl+C to return tea.Quit") + } +} + +func TestTUI_Model_EscClearsTextarea(t *testing.T) { + resolved := config.Resolved{ProfileName: "dev", Profile: config.Profile{DB: "mysql"}} + aiService := ai.NewService(config.AIConfig{}, nil) + m := NewModel(config.Options{}, resolved, aiService, "", false) + m.state = StateIdle + m.textarea.SetValue("draft prompt to clear") + + // Press Esc in StateIdle -> clears prompt + updated, _ := m.Update(tea.KeyMsg{Type: tea.KeyEsc}) + m = updated.(Model) + if m.textarea.Value() != "" { + t.Fatalf("expected textarea to be cleared after Esc, got %q", m.textarea.Value()) + } +} + +func TestTUI_Model_CtrlPProfileSwitching(t *testing.T) { + allProfiles := map[string]config.Profile{ + "dev": {DB: "mysql"}, + "prod": {DB: "pg"}, + } + resolved := config.Resolved{ + ProfileName: "dev", + Profile: allProfiles["dev"], + AllProfiles: allProfiles, + } + aiService := ai.NewService(config.AIConfig{}, nil) + m := NewModel(config.Options{}, resolved, aiService, "", false) + + // Press Ctrl+P -> switches profile to 'prod' + updated, cmd := m.Update(tea.KeyMsg{Type: tea.KeyCtrlP}) + m = updated.(Model) + if m.profileName != "prod" { + t.Fatalf("expected profileName to switch to 'prod', got %q", m.profileName) + } + if m.profile.DB != "pg" { + t.Fatalf("expected profile DB to be 'pg', got %q", m.profile.DB) + } + if cmd == nil { + t.Fatal("expected loadSchemaCmd after switching profile") + } +} + +func TestTUI_Model_ToolCallsAndTableRendering(t *testing.T) { + resolved := config.Resolved{ProfileName: "dev", Profile: config.Profile{DB: "mysql"}} + aiService := ai.NewService(config.AIConfig{}, nil) + m := NewModel(config.Options{}, resolved, aiService, "", false) + m.state = StateIdle + m.messages = append(m.messages, "", "") + + tc := ToolCallItem{ + ID: "tc_1", + Name: "execute_sql", + Summary: "SELECT * FROM users", + Detail: "SELECT * FROM users", + Result: "✓ Execution Success", + TableStateIndex: -1, + MsgIndex: 0, + IsExpanded: false, + } + m.toolCalls = append(m.toolCalls, tc) + + // Test focusToolCall & renderToolCall + m.focusToolCall(0) + if m.activeToolIdx != 0 { + t.Fatalf("expected activeToolIdx 0, got %d", m.activeToolIdx) + } + + // Toggle expanded via Ctrl+O + updated, _ := m.Update(tea.KeyMsg{Type: tea.KeyCtrlO}) + m = updated.(Model) + if !m.toolCalls[0].IsExpanded { + t.Fatal("expected toolCall to be expanded after Ctrl+O") + } + + // Test WindowSizeMsg + updated, _ = m.Update(tea.WindowSizeMsg{Width: 100, Height: 40}) + m = updated.(Model) + if m.width != 100 || m.height != 40 { + t.Fatalf("expected width 100, height 40, got %d, %d", m.width, m.height) + } +} + +func TestTUI_Model_ExportFlow(t *testing.T) { + resolved := config.Resolved{ProfileName: "dev", Profile: config.Profile{DB: "mysql"}} + aiService := ai.NewService(config.AIConfig{}, nil) + m := NewModel(config.Options{}, resolved, aiService, "", false) + + // Simulate aiResponseMsg returning TypeExport + updated, _ := m.Update(aiResponseMsg{ + response: &ai.AIResponse{ + Type: ai.TypeExport, + DatasetID: "res1", + Format: "csv", + FilePath: "test.csv", + }, + }) + m = updated.(Model) + if m.state != StateExportReady { + t.Fatalf("expected StateExportReady, got %v", m.state) + } + if m.pendingExport == nil { + t.Fatal("expected pendingExport to be non-nil") + } + + // Select option 3 (Deny / Cancel Export) via Key '3' + updated, cmd := m.Update(tea.KeyMsg{Type: tea.KeyRunes, Runes: []rune{'3'}}) + m = updated.(Model) + if m.state != StateThinking { + t.Fatalf("expected StateThinking after denying export, got %v", m.state) + } + if cmd == nil { + t.Fatal("expected runAgentStepCmd after export feedback") + } +} + +func TestTUI_Model_FullCoverage(t *testing.T) { + resolved := config.Resolved{ + ProfileName: "dev", + Profile: config.Profile{DB: "mysql", AllowPlaintext: true}, + } + aiService := ai.NewService(config.AIConfig{}, nil) + m := NewModel(config.Options{}, resolved, aiService, "", false) + + // 1. Test Init & loadSchemaCmd + initCmd := m.Init() + if initCmd == nil { + t.Fatal("expected non-nil Init Cmd") + } + + loadCmd := m.loadSchemaCmd() + if loadCmd == nil { + t.Fatal("expected non-nil loadSchemaCmd") + } + _ = loadCmd() // execute closure statements + + // 2. Test runAgentStepCmd & executeSQLCmd closures + stepCmd := m.runAgentStepCmd() + if stepCmd != nil { + _ = stepCmd() + } + + execCmd := m.executeSQLCmd("SELECT 1") + if execCmd != nil { + _ = execCmd() + } + + // 3. Test JS response and raw output rendering + updated, _ := m.Update(aiResponseMsg{ + response: &ai.AIResponse{ + Type: ai.TypeJS, + JSCode: "var x = 1;", + Explanation: "Run JS script", + }, + }) + m = updated.(Model) + + // 4. Test queryExecutedMsg with TableResult and TableState + res := &db.QueryResult{ + Columns: []string{"id", "name"}, + Rows: []map[string]any{{"id": 1, "name": "Alice"}}, + } + updated, _ = m.Update(queryExecutedMsg{ + result: res, + duration: 10 * time.Millisecond, + }) + m = updated.(Model) + + // 5. Test Key Navigation & renderTableState + m.tableStates = append(m.tableStates, TableState{ + Result: res, + MsgIndex: 0, + VerticalView: false, + }) + m.messages = []string{"msg0"} + m.renderTableState(0, true) + m.renderTableState(0, false) + m.tableStates[0].VerticalView = true + m.renderTableState(0, true) + + m.Update(tea.KeyMsg{Type: tea.KeyTab}) + m.Update(tea.KeyMsg{Type: tea.KeyRight}) + m.Update(tea.KeyMsg{Type: tea.KeyLeft}) + m.Update(tea.KeyMsg{Type: tea.KeyPgDown}) + m.Update(tea.KeyMsg{Type: tea.KeyPgUp}) + m.Update(tea.KeyMsg{Type: tea.KeyCtrlE}) + + // 6. Test Export Option 1 (Confirm Export) + m.state = StateExportReady + m.pendingExport = &PendingExport{ + DatasetID: "res1", + Format: "csv", + FilePath: "test.csv", + } + m.confirmOption = 0 + updated, _ = m.Update(tea.KeyMsg{Type: tea.KeyEnter}) + m = updated.(Model) + + // 7. Test Export Option 2 (Adjust Export Prompt) + m.state = StateExportReady + m.toolCalls = []ToolCallItem{{ID: "tc_1", Name: "export_data"}} + m.pendingExport = &PendingExport{ + DatasetID: "res1", + Format: "csv", + FilePath: "test.csv", + ToolIdx: 0, + } + m.confirmOption = 1 + updated, _ = m.Update(tea.KeyMsg{Type: tea.KeyEnter}) + m = updated.(Model) + if m.state != StateIdle { + t.Fatalf("expected StateIdle after Adjust Prompt option, got %v", m.state) + } +} + +func TestTUI_Model_ViewAndAllStatesCoverage(t *testing.T) { + resolved := config.Resolved{ + ProfileName: "dev", + Profile: config.Profile{DB: "mysql"}, + } + aiService := ai.NewService(config.AIConfig{}, nil) + m := NewModel(config.Options{}, resolved, aiService, "", true) // unsafeAllowWrite = true + m.autoExecute = true + + // 1. Test View in StateLoadingSchema + m.state = StateLoadingSchema + if view := m.View(); !strings.Contains(view, "READ-WRITE") || !strings.Contains(view, "AUTO-EXEC") { + t.Fatalf("expected badges in header view, got:\n%s", view) + } + + // 2. Test View in all states + states := []State{ + StateIdle, StateThinking, StateSQLReady, StateExecuting, StateExportReady, + } + for _, st := range states { + m.state = st + _ = m.View() + } + + // 3. Test schemaLoadedMsg with error + updated, _ := m.Update(schemaLoadedMsg{ + err: errors.New("XSQL_CFG_INVALID", "invalid config", nil), + }) + m = updated.(Model) + + // 4. Test aiResponseMsg with error & TypeText + updated, _ = m.Update(aiResponseMsg{ + err: errors.New("XSQL_AI_API_ERROR", "api failed", nil), + }) + m = updated.(Model) + if m.state != StateIdle { + t.Fatalf("expected StateIdle after AI error, got %v", m.state) + } + + updated, _ = m.Update(aiResponseMsg{ + response: &ai.AIResponse{ + Type: ai.TypeText, + Explanation: "Here is text response", + }, + }) + m = updated.(Model) + if m.state != StateIdle { + t.Fatalf("expected StateIdle after TypeText response, got %v", m.state) + } + + // 5. Test queryExecutedMsg with error + updated, _ = m.Update(queryExecutedMsg{ + err: errors.New("XSQL_SQL_SYNTAX_ERROR", "syntax error", nil), + }) + m = updated.(Model) + + // 6. Test Option keys '1', '2', '3' in StateSQLReady + m.state = StateSQLReady + m.currentSQL = "SELECT 1;" + + m.Update(tea.KeyMsg{Type: tea.KeyRunes, Runes: []rune{'2'}}) + m.state = StateSQLReady + m.Update(tea.KeyMsg{Type: tea.KeyRunes, Runes: []rune{'3'}}) + m.state = StateSQLReady + m.Update(tea.KeyMsg{Type: tea.KeyRunes, Runes: []rune{'1'}}) + + // 7. Test focusToolCall out of bounds & renderTableState invalid indices + m.focusToolCall(-1) + m.focusToolCall(999) + m.renderTableState(-1, false) + m.renderTableState(999, false) +} diff --git a/internal/tui/render.go b/internal/tui/render.go new file mode 100644 index 0000000..37a384d --- /dev/null +++ b/internal/tui/render.go @@ -0,0 +1,93 @@ +package tui + +import ( + "strings" + + "github.com/alecthomas/chroma/v2" + "github.com/alecthomas/chroma/v2/lexers" + "github.com/charmbracelet/glamour" + "github.com/charmbracelet/lipgloss" +) + +var ( + // High-Contrast Adaptive Styles for SQL & JS Syntax Highlighting (Light & Dark mode compatible) + KeywordStyle = lipgloss.NewStyle().Bold(true).Foreground(lipgloss.AdaptiveColor{Light: "#6D28D9", Dark: "#C084FC"}) // Rich Purple + StringStyle = lipgloss.NewStyle().Foreground(lipgloss.AdaptiveColor{Light: "#047857", Dark: "#34D399"}) // Emerald Green + NumberStyle = lipgloss.NewStyle().Foreground(lipgloss.AdaptiveColor{Light: "#B45309", Dark: "#FBBF24"}) // Amber Gold + NameStyle = lipgloss.NewStyle().Bold(true).Foreground(lipgloss.AdaptiveColor{Light: "#0369A1", Dark: "#38BDF8"}) // Sky Blue + CommentStyle = lipgloss.NewStyle().Italic(true).Foreground(lipgloss.AdaptiveColor{Light: "#64748B", Dark: "#94A3B8"}) // Slate Grey + DefaultTxtStyle = lipgloss.NewStyle().Bold(true).Foreground(lipgloss.AdaptiveColor{Light: "#0F172A", Dark: "#F8FAFC"}) // Deep Slate / Crisp White +) + +// RenderMarkdown renders markdown text using Glamour with rich ANSI terminal styling. +func RenderMarkdown(md string, width int) string { + md = strings.TrimSpace(md) + if md == "" { + return "" + } + if width <= 10 { + width = 80 + } + r, err := glamour.NewTermRenderer( + glamour.WithStandardStyle("auto"), + glamour.WithWordWrap(width-6), + ) + if err != nil { + return md + } + out, err := r.Render(md) + if err != nil { + return md + } + return strings.TrimSpace(out) +} + +// HighlightCode renders adaptive, high-contrast syntax-highlighted code for light & dark terminals. +func HighlightCode(code string, lexerName string) string { + code = strings.TrimSpace(code) + if code == "" { + return "" + } + + lexer := lexers.Get(lexerName) + if lexer == nil { + lexer = lexers.Fallback + } + lexer = chroma.Coalesce(lexer) + + iterator, err := lexer.Tokenise(nil, code) + if err != nil { + return DefaultTxtStyle.Render(code) + } + + var sb strings.Builder + for _, t := range iterator.Tokens() { + val := t.Value + switch t.Type { + case chroma.Keyword, chroma.KeywordReserved, chroma.KeywordType, chroma.KeywordNamespace: + sb.WriteString(KeywordStyle.Render(val)) + case chroma.String, chroma.StringChar, chroma.StringSingle, chroma.StringDouble, chroma.StringBacktick: + sb.WriteString(StringStyle.Render(val)) + case chroma.Number, chroma.NumberInteger, chroma.NumberFloat, chroma.NumberHex, chroma.NumberOct: + sb.WriteString(NumberStyle.Render(val)) + case chroma.Name, chroma.NameAttribute, chroma.NameClass, chroma.NameFunction, chroma.NameTag: + sb.WriteString(NameStyle.Render(val)) + case chroma.Comment, chroma.CommentSingle, chroma.CommentMultiline: + sb.WriteString(CommentStyle.Render(val)) + default: + sb.WriteString(DefaultTxtStyle.Render(val)) + } + } + + return sb.String() +} + +// HighlightSQL applies adaptive high-contrast syntax highlighting to SQL statements. +func HighlightSQL(sqlStr string) string { + return HighlightCode(sqlStr, "sql") +} + +// HighlightJS applies adaptive high-contrast syntax highlighting to JavaScript code blocks. +func HighlightJS(jsStr string) string { + return HighlightCode(jsStr, "javascript") +} diff --git a/internal/tui/render_test.go b/internal/tui/render_test.go new file mode 100644 index 0000000..8d5b3da --- /dev/null +++ b/internal/tui/render_test.go @@ -0,0 +1,45 @@ +package tui + +import ( + "strings" + "testing" +) + +func TestRenderMarkdown(t *testing.T) { + out := RenderMarkdown("# Hello World\nThis is a **test**.", 80) + if !strings.Contains(out, "Hello World") { + t.Errorf("expected rendered markdown to contain 'Hello World', got %q", out) + } + + empty := RenderMarkdown("", 80) + if empty != "" { + t.Errorf("expected empty string for empty input, got %q", empty) + } + + narrow := RenderMarkdown("# Header", 5) + if !strings.Contains(narrow, "Header") { + t.Errorf("expected narrow width fallback, got %q", narrow) + } +} + +func TestHighlightCode(t *testing.T) { + sqlOut := HighlightSQL("SELECT * FROM users WHERE id = 1;") + if sqlOut == "" || !strings.Contains(sqlOut, "SELECT") { + t.Errorf("expected highlighted SQL, got %q", sqlOut) + } + + jsOut := HighlightJS("var data = res1;\nconsole.log(data);") + if jsOut == "" || !strings.Contains(jsOut, "data") { + t.Errorf("expected highlighted JS, got %q", jsOut) + } + + emptyCode := HighlightCode("", "sql") + if emptyCode != "" { + t.Errorf("expected empty string for empty code, got %q", emptyCode) + } + + fallbackCode := HighlightCode("some raw string", "unknown_lexer") + if fallbackCode == "" { + t.Errorf("expected fallback highlight for unknown lexer") + } +} diff --git a/internal/tui/styles.go b/internal/tui/styles.go new file mode 100644 index 0000000..e941cea --- /dev/null +++ b/internal/tui/styles.go @@ -0,0 +1,146 @@ +package tui + +import "github.com/charmbracelet/lipgloss" + +var ( + // Palette Colors (Adaptive Catppuccin / Tokyo Night Theme) + PrimaryColor = lipgloss.AdaptiveColor{Light: "#6D28D9", Dark: "#A78BFA"} + SecondaryColor = lipgloss.AdaptiveColor{Light: "#059669", Dark: "#34D399"} + AccentColor = lipgloss.AdaptiveColor{Light: "#BE185D", Dark: "#F472B6"} + WarningColor = lipgloss.AdaptiveColor{Light: "#D97706", Dark: "#FBBF24"} + ErrorColor = lipgloss.AdaptiveColor{Light: "#E11D48", Dark: "#F87171"} + InfoColor = lipgloss.AdaptiveColor{Light: "#0284C7", Dark: "#38BDF8"} + MutedColor = lipgloss.AdaptiveColor{Light: "#64748B", Dark: "#94A3B8"} + TextNormal = lipgloss.AdaptiveColor{Light: "#0F172A", Dark: "#F8FAFC"} + HeaderBg = lipgloss.AdaptiveColor{Light: "#E2E8F0", Dark: "#1E293B"} + + // Header Container & Badges + HeaderBarStyle = lipgloss.NewStyle(). + Background(HeaderBg). + Padding(0, 1) + + HeaderTitleBadge = lipgloss.NewStyle(). + Bold(true). + Foreground(lipgloss.Color("#FFFFFF")). + Background(PrimaryColor). + Padding(0, 1) + + HeaderProfileBadge = lipgloss.NewStyle(). + Bold(true). + Foreground(lipgloss.Color("#FFFFFF")). + Background(InfoColor). + Padding(0, 1) + + BadgeReadOnly = lipgloss.NewStyle(). + Bold(true). + Foreground(lipgloss.Color("#FFFFFF")). + Background(SecondaryColor). + Padding(0, 1) + + BadgeReadWrite = lipgloss.NewStyle(). + Bold(true). + Foreground(lipgloss.Color("#FFFFFF")). + Background(WarningColor). + Padding(0, 1) + + BadgeAutoExec = lipgloss.NewStyle(). + Bold(true). + Foreground(lipgloss.Color("#FFFFFF")). + Background(AccentColor). + Padding(0, 1) + + BadgeManualApprove = lipgloss.NewStyle(). + Bold(true). + Foreground(lipgloss.AdaptiveColor{Light: "#0F172A", Dark: "#E2E8F0"}). + Background(lipgloss.AdaptiveColor{Light: "#CBD5E1", Dark: "#475569"}). + Padding(0, 1) + + // Keybinding Badges + KeyBadgeStyle = lipgloss.NewStyle(). + Bold(true). + Foreground(lipgloss.AdaptiveColor{Light: "#0F172A", Dark: "#F8FAFC"}). + Background(lipgloss.AdaptiveColor{Light: "#CBD5E1", Dark: "#334155"}). + Padding(0, 1) + + KeyLabelStyle = lipgloss.NewStyle(). + Foreground(MutedColor) + + // SQL Preview Box + SQLBoxStyle = lipgloss.NewStyle(). + Border(lipgloss.RoundedBorder()). + BorderForeground(PrimaryColor). + Padding(0, 1). + MarginTop(1). + MarginBottom(1) + + SQLTitleStyle = lipgloss.NewStyle(). + Bold(true). + Foreground(AccentColor) + + SQLCodeStyle = lipgloss.NewStyle(). + Bold(true). + Foreground(InfoColor) + + // User & AI Message Tags + UserTagStyle = lipgloss.NewStyle(). + Bold(true). + Foreground(lipgloss.Color("#FFFFFF")). + Background(SecondaryColor). + Padding(0, 1) + + AITagStyle = lipgloss.NewStyle(). + Bold(true). + Foreground(lipgloss.Color("#FFFFFF")). + Background(PrimaryColor). + Padding(0, 1) + + ExecutingTagStyle = lipgloss.NewStyle(). + Bold(true). + Foreground(lipgloss.Color("#FFFFFF")). + Background(WarningColor). + Padding(0, 1) + + SuccessBadgeStyle = lipgloss.NewStyle(). + Bold(true). + Foreground(SecondaryColor) + + WarningBadgeStyle = lipgloss.NewStyle(). + Bold(true). + Foreground(WarningColor) + + MetricsStyle = lipgloss.NewStyle(). + Foreground(MutedColor). + Italic(true) + + AIResponseStyle = lipgloss.NewStyle(). + Foreground(TextNormal). + PaddingLeft(1) + + ErrorMsgStyle = lipgloss.NewStyle(). + Bold(true). + Foreground(ErrorColor). + PaddingLeft(1) + + PromptPrefixStyle = lipgloss.NewStyle(). + Bold(true). + Foreground(PrimaryColor) + + // Collapsible Tool Call Badges + ToolCollapsedBadge = lipgloss.NewStyle(). + Bold(true). + Foreground(lipgloss.AdaptiveColor{Light: "#475569", Dark: "#94A3B8"}). + Background(lipgloss.AdaptiveColor{Light: "#E2E8F0", Dark: "#1E293B"}). + Padding(0, 1) + + ToolExpandedBadge = lipgloss.NewStyle(). + Bold(true). + Foreground(lipgloss.Color("#FFFFFF")). + Background(WarningColor). + Padding(0, 1) + + ToolDetailStyle = lipgloss.NewStyle(). + Border(lipgloss.NormalBorder(), false, false, false, true). + BorderForeground(MutedColor). + PaddingLeft(1). + Foreground(MutedColor) +) diff --git a/tests/e2e/ai_test.go b/tests/e2e/ai_test.go new file mode 100644 index 0000000..13f9a42 --- /dev/null +++ b/tests/e2e/ai_test.go @@ -0,0 +1,198 @@ +//go:build e2e + +package e2e + +import ( + "bytes" + "context" + "fmt" + "io" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "strings" + "testing" + "time" + + tea "github.com/charmbracelet/bubbletea" + + "github.com/zx06/xsql/internal/ai" + "github.com/zx06/xsql/internal/config" + "github.com/zx06/xsql/internal/db" + "github.com/zx06/xsql/internal/tui" +) + +func TestE2E_AI_Service_With_MockOpenAI(t *testing.T) { + // 1. Setup Mock OpenAI Server with Tool Call response + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/chat/completions" { + http.Error(w, "not found", http.StatusNotFound) + return + } + if r.Header.Get("Authorization") != "Bearer test-e2e-key" { + http.Error(w, "unauthorized", http.StatusUnauthorized) + return + } + + bodyBytes, _ := io.ReadAll(r.Body) + bodyStr := string(bodyBytes) + + // Assert request contains schema context + if !strings.Contains(bodyStr, "DATABASE SCHEMA") { + t.Errorf("request body missing schema context: %s", bodyStr) + } + + respBody := `{ + "id": "chatcmpl-e2e-123", + "object": "chat.completion", + "created": 1677652288, + "model": "gpt-4o", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": null, + "tool_calls": [ + { + "id": "call_e2e_123", + "type": "function", + "function": { + "name": "execute_sql", + "arguments": "{\"sql\":\"SELECT COUNT(*) FROM users;\",\"explanation\":\"Returns total user count.\"}" + } + } + ] + }, + "finish_reason": "tool_calls" + } + ] + }` + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(respBody)) + })) + defer server.Close() + + // 2. Setup AI Config and Service + aiCfg := config.AIConfig{ + Provider: "openai", + BaseURL: server.URL, + APIKey: "test-e2e-key", + Model: "gpt-4o", + MaxTokens: 2048, + } + + client := ai.NewClient(aiCfg, server.Client()) + service := ai.NewService(aiCfg, client) + + mockSchema := &db.SchemaInfo{ + Database: "e2e_db", + Tables: []db.Table{ + { + Name: "users", + Columns: []db.Column{ + {Name: "id", Type: "bigint", PrimaryKey: true}, + }, + }, + }, + } + + // 3. Test GenerateSQL + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + sqlResp, xe := service.GenerateSQL(ctx, "How many users exist?", mockSchema, "mysql") + if xe != nil { + t.Fatalf("unexpected error generating SQL: %v", xe) + } + + if sqlResp.SQL != "SELECT COUNT(*) FROM users;" { + t.Errorf("expected SQL 'SELECT COUNT(*) FROM users;', got %q", sqlResp.SQL) + } + if sqlResp.Explanation != "Returns total user count." { + t.Errorf("expected explanation 'Returns total user count.', got %q", sqlResp.Explanation) + } +} + +func TestE2E_AI_TUI_Terminal_Pipe(t *testing.T) { + // Setup Mock Server for AI BaseURL + aiServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + respBody := `{ + "id": "chatcmpl-tui", + "object": "chat.completion", + "created": 1677652288, + "model": "test-model", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": null, + "tool_calls": [ + { + "id": "call_tui", + "type": "function", + "function": { + "name": "execute_sql", + "arguments": "{\"sql\":\"SELECT COUNT(*) FROM users;\",\"explanation\":\"Returns total user count.\"}" + } + } + ] + }, + "finish_reason": "tool_calls" + } + ] + }` + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(respBody)) + })) + defer aiServer.Close() + + // Setup temporary xsql config file + tmpDir := t.TempDir() + cfgPath := filepath.Join(tmpDir, "xsql.yaml") + cfgContent := fmt.Sprintf(`profiles: + dev: + db: mysql + host: 127.0.0.1 + port: 3306 + user: root + database: test +ai: + base_url: "%s" + model: "test-model" + api_key: "test-key" +`, aiServer.URL) + if err := os.WriteFile(cfgPath, []byte(cfgContent), 0600); err != nil { + t.Fatal(err) + } + + resolved, xe := config.Resolve(config.Options{ConfigPath: cfgPath, CLIProfile: "dev", CLIProfileSet: true}) + if xe != nil { + t.Fatalf("failed to resolve config: %v", xe) + } + + aiClient := ai.NewClient(resolved.AI, aiServer.Client()) + aiService := ai.NewService(resolved.AI, aiClient) + model := tui.NewModel(config.Options{}, resolved, aiService, "Show total users", false) + + inBuf := bytes.NewBufferString("\n") // Press Enter + outBuf := &bytes.Buffer{} + + p := tea.NewProgram(model, tea.WithInput(inBuf), tea.WithOutput(outBuf)) + + go func() { + time.Sleep(300 * time.Millisecond) + p.Quit() + }() + + if _, err := p.Run(); err != nil { + t.Fatalf("TUI program run failed: %v", err) + } + + outputStr := outBuf.String() + if !strings.Contains(outputStr, "xsql AI") { + t.Errorf("expected TUI terminal output to contain header 'xsql AI', got:\n%s", outputStr) + } +} diff --git a/tests/e2e/proxy_test.go b/tests/e2e/proxy_test.go index a720fd0..d7a169f 100644 --- a/tests/e2e/proxy_test.go +++ b/tests/e2e/proxy_test.go @@ -356,18 +356,19 @@ func TestProxy_Help(t *testing.T) { } func TestProxy_MissingProfileFlag(t *testing.T) { - // Test that proxy fails when -p flag is not provided - stdout, _, exitCode := runXSQL(t, "proxy", "--format", "json") + // Test that proxy fails when profile is missing or invalid + emptyCfg := createTempConfig(t, "profiles: {}") + stdout, _, exitCode := runXSQL(t, "proxy", "--config", emptyCfg, "--format", "json") // Should fail with config error if exitCode == 0 { - t.Error("expected non-zero exit code when -p flag is not provided") + t.Error("expected non-zero exit code when profile is not provided") } var resp Response if err := json.Unmarshal([]byte(stdout), &resp); err == nil { if resp.OK { - t.Error("expected ok=false when -p flag is missing") + t.Error("expected ok=false when profile is missing") } } }