diff --git a/cmd/godex/main.go b/cmd/godex/main.go index 69ffc76..1bd3a05 100644 --- a/cmd/godex/main.go +++ b/cmd/godex/main.go @@ -35,6 +35,7 @@ import ( "github.com/cheikh2shift/godex/internal/mcp" "github.com/cheikh2shift/godex/internal/ml" "github.com/cheikh2shift/godex/internal/providers" + "github.com/cheikh2shift/godex/internal/scheduler" "github.com/cheikh2shift/godex/internal/wizard" "github.com/cheikh2shift/godex/modelquery" ) @@ -62,7 +63,7 @@ type MCPServer interface { Close() error } -var slashCommands = []string{"/add-path ", "/remove-path ", "/paths", "/tools", "/clear-context", "/commit ", "/commit-pull ", "/commit-merge ", "/commit-search ", "/exit", "/quit", "/q", "/save", "/save-exit", "/kill ", "/killbg", "/bg", "/clear", "/help", "/model", "/model-persist"} +var slashCommands = []string{"/add-path ", "/remove-path ", "/paths", "/tools", "/clear-context", "/commit ", "/commit-pull ", "/commit-merge ", "/commit-search ", "/exit", "/quit", "/q", "/save", "/save-exit", "/kill ", "/killbg", "/bg", "/clear", "/help", "/model", "/model-persist", "/schedule", "/schedule-clear"} var ( greenOrb = lipgloss.NewStyle().Foreground(lipgloss.Color("82")).Render("●") @@ -73,6 +74,7 @@ var ( pathPromptMu sync.Mutex hivePendingStop chan struct{} summarisingStart chan struct{} + sched *scheduler.Scheduler ) var thinkingStyle = lipgloss.NewStyle(). @@ -319,10 +321,8 @@ func main() { setter.SetStatusChannel(statusCh) } - fmt.Println() - if prompt != "" { - err := runSinglePrompt(ctx, provider, prompt, autoConfirm, debugMode, servers) + err = runSinglePrompt(ctx, provider, prompt, autoConfirm, debugMode, servers) cleanup(servers) if err != nil { log.Fatalf("prompt failed: %v", err) @@ -330,9 +330,58 @@ func main() { return } + sched, err = scheduler.NewDefault() + if err != nil { + log.Printf("[Scheduler] Failed to initialize: %v", err) + } else { + sched.SetStatusChannel(statusCh) + if debugMode { + fmt.Println(muted.Render("[Scheduler] Initialized with statusCh")) + } + sched.SetOnTaskFinished(func(*scheduler.ScheduledTask) { + playSound() + }) + sched.SetProviderGetter(&schedulerProviderGetter{provider: provider}) + var serverSlice []interface{} = make([]interface{}, len(servers)) + for i, s := range servers { + serverSlice[i] = s + } + sched.SetServers(serverSlice) + + schedulerServer := mcp.NewSchedulerServer(sched, &schedulerProviderGetter{provider: provider}) + servers = append(servers, schedulerServer) + for _, tool := range schedulerServer.Tools() { + var inputSchema map[string]interface{} + if err := json.Unmarshal(tool.InputSchema, &inputSchema); err != nil { + inputSchema = map[string]interface{}{} + } + providerTools = append(providerTools, providers.Tool{ + Name: tool.Name, + Description: tool.Description, + InputSchema: inputSchema, + }) + } + agent.SetProviderTools(provider, providerTools) + if debugMode { + fmt.Println(muted.Render("[Scheduler] Initialized")) + } + } + // Get working directory for session files wd, _ := os.Getwd() + if sched != nil { + maxRounds := 10 + if provider.MaxToolRounds != nil && *provider.MaxToolRounds > 0 { + maxRounds = *provider.MaxToolRounds + } + toolTimeout := 180 + if provider.ToolTimeout != nil && *provider.ToolTimeout > 0 { + toolTimeout = *provider.ToolTimeout + } + sched.SetConfig(maxRounds, toolTimeout, wd, runtime.GOOS) + } + // Initialize history database (stored in tmp directory) var historyDB *history.HistoryDB historyDB, err = history.NewDefault() @@ -852,6 +901,14 @@ promptLoop: } continue } + if input == "/schedule-clear" { + handleSchedule("/schedule clear") + continue + } + if strings.HasPrefix(input, "/schedule") { + handleSchedule(input) + continue + } nativeToolCalls := llmProvider.SupportsNativeToolCalls() var toolsSection string @@ -955,12 +1012,14 @@ promptLoop: fmt.Printf("[Round %d/%d] Got %d tool calls, isToolCallResponse=%v\n", round+1, maxToolRounds, len(toolCalls), isToolCallResponse) - if len(toolCalls) > 0 && isToolCallResponse { + totalToolCalls := len(toolCalls) + len(missingTools) + if totalToolCalls > 0 && isToolCallResponse { prevNoTool = false } + if !isToolCallResponse || len(toolCalls) == 0 { - if !looksLikeToolCall(resp) { + if !looksLikeToolCall(resp) && round > 0 { go playSound() fmt.Printf("\n\n%s\n", renderMarkdown(resp)) break @@ -1036,12 +1095,12 @@ promptLoop: } } - toolDesc := getToolDescription(servers, toolName) + /*toolDesc := getToolDescription(servers, toolName) if toolDesc != "" { fmt.Print("\033[90m") fmt.Printf(" %s\n", toolDesc) fmt.Print("\033[0m") - } + }*/ var argsStr string for k, v := range args { @@ -1234,6 +1293,9 @@ promptLoop: if historyDB != nil { historyDB.Close() } + if sched != nil { + sched.Close() + } fmt.Println("Goodbye!") } @@ -1803,17 +1865,6 @@ func handleKill(servers []MCPServer, input string) { fmt.Println("No bash server found") } -func getToolDescription(servers []MCPServer, toolName string) string { - for _, server := range servers { - for _, tool := range server.Tools() { - if tool.Name == toolName { - return tool.Description - } - } - } - return "" -} - func handleListBg(servers []MCPServer) { for _, server := range servers { if bs, ok := server.(*mcp.BashServer); ok { @@ -1955,6 +2006,7 @@ Commands: /commit-merge - Merge committed history into current state /commit-search - Search commits by ref or message /clear-context - Reset context counter and LLM client + /schedule - List scheduled tasks (see /schedule help for more) /save, /save-exit - Save session and exit /kill - Kill a background process by PID /killbg - Kill all background processes @@ -2192,7 +2244,7 @@ func runToolLoop(ctx context.Context, provider *config.Provider, servers []MCPSe if !isToolCallResponse || len(toolCalls) == 0 { if !hasFinal { - if !looksLikeToolCall(resp) { + if !looksLikeToolCall(resp) && round > 0 { if nocont { fmt.Printf("\n\n%s\n", renderMarkdown(resp)) } diff --git a/cmd/godex/prompt.go b/cmd/godex/prompt.go index 568c4f5..1ac38ab 100644 --- a/cmd/godex/prompt.go +++ b/cmd/godex/prompt.go @@ -17,6 +17,7 @@ import ( "github.com/charmbracelet/lipgloss" "github.com/cheikh2shift/godex/internal/history" "github.com/cheikh2shift/godex/internal/hive" + "github.com/cheikh2shift/godex/internal/scheduler" ) var ErrPromptAborted = errors.New("prompt aborted") @@ -74,6 +75,7 @@ type promptModel struct { searchResults []string statusMessage string commitList []string + scheduleList []string statusCh <-chan string delegateCh <-chan hive.HiveStats delegateCount int @@ -166,9 +168,10 @@ func (m promptModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) { if m.statusHidden { return m, nil } - if (m.statusMessage != "" || len(m.commitList) > 0) && msg.Type != tea.KeyTab { + if (m.statusMessage != "" || len(m.commitList) > 0 || len(m.scheduleList) > 0) && msg.Type != tea.KeyTab { m.statusMessage = "" m.commitList = nil + m.scheduleList = nil } if msg.Type == tea.KeyCtrlR { if m.searchMode { @@ -254,6 +257,12 @@ func (m promptModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) { m.applyStatusMessage(msg) return m, waitStatusCmd(m.statusCh) } + if _, ok := msg.(statusMsgPoll); ok { + return m, waitStatusCmd(m.statusCh) + } + if _, ok := msg.(delegateMsgPoll); ok { + return m, waitDelegateCmd(m.delegateCh) + } if msg, ok := msg.(delegateMsg); ok { m.delegateCount = msg.Count m.delegateLatest = msg.Latest @@ -409,6 +418,14 @@ func (m promptModel) View() string { b.WriteByte('\n') } } + if len(m.scheduleList) > 0 { + b.WriteByte('\n') + listStyle := lipgloss.NewStyle().Foreground(lipgloss.Color("75")) + for _, id := range m.scheduleList { + b.WriteString(listStyle.Render(id)) + b.WriteByte('\n') + } + } if m.statusMessage != "" { b.WriteByte('\n') b.WriteString(lipgloss.NewStyle().Foreground(lipgloss.Color("240")).Render(m.statusMessage)) @@ -550,6 +567,11 @@ func (m *promptModel) handleTab() { return } } + if strings.HasPrefix(value, "/schedule ") && strings.Contains(value, " ") { + if m.handleScheduleTab(value) { + return + } + } if value != "" && !strings.HasPrefix(value, "/") { return } @@ -589,6 +611,62 @@ func (m *promptModel) handleTab() { m.showCompletions = true } +func (m *promptModel) handleScheduleTab(value string) bool { + hasTrailingSpace := strings.HasSuffix(value, " ") + parts := strings.Fields(value) + if hasTrailingSpace { + parts = append(parts, "") + } + + var ( + prefix string + nextPrefix string + ) + + if len(parts) == 2 { + prefix = strings.ToUpper(parts[1]) + nextPrefix = "/schedule " + } else if len(parts) == 3 && (parts[1] == "stop" || parts[1] == "remove") { + prefix = strings.ToUpper(parts[2]) + nextPrefix = "/schedule " + parts[1] + " " + } else { + return false + } + + if sched == nil { + return false + } + tasks := sched.ListTasks() + if len(tasks) == 0 { + return false + } + var matches []string + for _, t := range tasks { + task, ok := t.(*scheduler.ScheduledTask) + if !ok { + continue + } + if strings.HasPrefix(task.ID, prefix) { + matches = append(matches, task.ID) + } + } + if len(matches) == 0 { + return false + } + if len(matches) == 1 { + m.input.SetValue(nextPrefix + matches[0]) + m.input.SetCursor(len([]rune(m.input.Value()))) + m.resetCompletion() + return true + } + m.scheduleList = nil + for _, id := range matches { + m.scheduleList = append(m.scheduleList, id) + } + m.resetCompletion() + return true +} + func (m *promptModel) handleCommitPullTab(value string) bool { if m.historyDB == nil { return false @@ -863,27 +941,50 @@ func waitStatusCmd(ch <-chan string) tea.Cmd { return nil } return func() tea.Msg { - msg, ok := <-ch - if !ok { - return nil + select { + case msg, ok := <-ch: + if !ok { + return nil + } + //fmt.Printf("DEBUG waitStatusCmd received: %s\n", msg) + return statusMsg(msg) + case <-time.After(50 * time.Millisecond): + return statusMsgPoll{} } - return statusMsg(msg) } } +type statusMsgPoll struct{} + +func (statusMsgPoll) Init() {} +func (p statusMsgPoll) Update(msg tea.Msg) (tea.Msg, tea.Cmd) { + return msg, nil +} + func waitDelegateCmd(ch <-chan hive.HiveStats) tea.Cmd { if ch == nil { return nil } return func() tea.Msg { - msg, ok := <-ch - if !ok { - return nil + select { + case msg, ok := <-ch: + if !ok { + return nil + } + return delegateMsg{Count: msg.DelegatedCount, Latest: msg.LatestCommand} + case <-time.After(50 * time.Millisecond): + return delegateMsgPoll{} } - return delegateMsg{Count: msg.DelegatedCount, Latest: msg.LatestCommand} } } +type delegateMsgPoll struct{} + +func (delegateMsgPoll) Init() {} +func (p delegateMsgPoll) Update(msg tea.Msg) (tea.Msg, tea.Cmd) { + return msg, nil +} + type commitRow struct { primary string secondary string diff --git a/cmd/godex/scheduler.go b/cmd/godex/scheduler.go new file mode 100644 index 0000000..e2e4fbb --- /dev/null +++ b/cmd/godex/scheduler.go @@ -0,0 +1,230 @@ +package main + +import ( + "fmt" + "strings" + + "github.com/cheikh2shift/godex/internal/agent" + "github.com/cheikh2shift/godex/internal/config" + "github.com/cheikh2shift/godex/internal/scheduler" +) + +type schedulerProviderGetter struct { + provider *config.Provider +} + +func (s *schedulerProviderGetter) GetProvider(cfg interface{}) (interface{}, error) { + cfgMap, ok := cfg.(map[string]interface{}) + if !ok { + return nil, fmt.Errorf("invalid config type") + } + + providerType, _ := cfgMap["type"].(string) + providerName, _ := cfgMap["name"].(string) + providerModel, _ := cfgMap["model"].(string) + + // The scheduler often requests a "current" provider. For providers like OpenRouter, + // the API key lives on the full config.Provider struct, so we must pass it through. + useCurrent := s.provider != nil && (providerType == "" || providerType == "unknown" || providerType == "current" || providerName == "" || providerName == "current") + if useCurrent { + // Reinitialize to ensure runtime auth (OpenRouter keys, env) is present for scheduler runs. + // NOTE: We intentionally do not override the model here to avoid mutating global config. + _ = providerModel + return agent.ReinitializeProvider(s.provider) + } + + return agent.GetProviderFromConfig(&agent.ProviderConfig{ + Type: providerType, + Name: providerName, + Model: providerModel, + }) +} + +func handleSchedule(input string) { + parts := strings.Fields(input) + if len(parts) >= 2 && (parts[1] == "help" || parts[1] == "-h" || parts[1] == "--help") { + fmt.Println("Scheduler - Schedule prompts to run at specific times or intervals") + fmt.Println("") + fmt.Println("Usage:") + fmt.Println(" /schedule - List all scheduled tasks") + fmt.Println(" /schedule - Show last output of a task") + fmt.Println(" /schedule stop - Stop a running task by 4-char ID") + fmt.Println(" /schedule stop --all - Stop all running tasks") + fmt.Println(" /schedule remove - Remove a task from the schedule by 4-char ID") + fmt.Println(" /schedule remove --all - Remove all tasks from the schedule") + fmt.Println(" /schedule clear - Stop and remove all tasks") + fmt.Println(" /schedule help - Show this help message") + fmt.Println("") + fmt.Println("LLM Tool (scheduler):") + fmt.Println(" prompt: The prompt to execute") + fmt.Println(" interval_sec: Repeating interval in seconds (e.g., 60, 3600)") + fmt.Println(" run_at: Time in HH:MM format (e.g., '14:30') or 'now' for immediate execution") + fmt.Println("") + fmt.Println("Examples:") + fmt.Println(" scheduler(prompt='Summarize my emails', interval_sec=3600)") + fmt.Println(" scheduler(prompt='Check system status', run_at='09:00')") + fmt.Println(" scheduler(prompt='Generate report', run_at='now')") + return + } + + if sched == nil { + fmt.Println("[Scheduler] Not initialized") + return + } + + if len(parts) == 1 { + listScheduledTasks() + return + } + + if len(parts) == 2 && (parts[1] == "clear" || parts[1] == "clear-all") { + tasks := sched.ListTasks() + if len(tasks) == 0 { + fmt.Println("No scheduled tasks") + return + } + sched.StopAllTasks() + removed := 0 + for _, t := range tasks { + task, ok := t.(*scheduler.ScheduledTask) + if !ok || task == nil { + continue + } + if sched.RemoveTask(task.ID) { + removed++ + } + } + fmt.Printf("Cleared %d task(s)\n", removed) + return + } + + if len(parts) >= 3 && parts[1] == "stop" { + if parts[2] == "all" || parts[2] == "--all" || parts[2] == "*" { + sched.StopAllTasks() + fmt.Println("Stopped all running tasks") + return + } + id := strings.ToUpper(parts[2]) + if len(id) != 4 { + fmt.Println("Invalid task ID. Use 4-character identifier (e.g., ABC1)") + return + } + if sched.StopTask(id) { + fmt.Printf("Task %s stopped\n", id) + } else { + fmt.Printf("Task %s not found\n", id) + } + return + } + + if len(parts) >= 3 && parts[1] == "remove" { + if parts[2] == "all" || parts[2] == "--all" || parts[2] == "*" { + tasks := sched.ListTasks() + if len(tasks) == 0 { + fmt.Println("No scheduled tasks") + return + } + removed := 0 + for _, t := range tasks { + task, ok := t.(*scheduler.ScheduledTask) + if !ok || task == nil { + continue + } + if sched.RemoveTask(task.ID) { + removed++ + } + } + fmt.Printf("Removed %d task(s)\n", removed) + return + } + id := strings.ToUpper(parts[2]) + if len(id) != 4 { + fmt.Println("Invalid task ID. Use 4-character identifier (e.g., ABC1)") + return + } + if sched.RemoveTask(id) { + fmt.Printf("Task %s removed\n", id) + } else { + fmt.Printf("Task %s not found\n", id) + } + return + } + + if len(parts) == 2 && len(parts[1]) == 4 { + id := strings.ToUpper(parts[1]) + task := sched.GetTask(id) + if task == nil { + fmt.Printf("Task %s not found\n", id) + return + } + st, ok := task.(*scheduler.ScheduledTask) + if !ok || st == nil { + fmt.Printf("Task %s not found\n", id) + return + } + if st.LastOutput == "" { + + if st.LastError != "" { + fmt.Printf("Error for task %s:\n%s\n", id, st.LastError) + return + } + fmt.Printf("No output for task %s\n", id) + return + } + fmt.Printf("Output for task %s:\n%s\n", id, renderMarkdown(st.LastOutput)) + return + } + + fmt.Println("Usage:") + fmt.Println(" /schedule - List all scheduled tasks") + fmt.Println(" /schedule - Show last output of a task") + fmt.Println(" /schedule stop - Stop a task by 4-char ID") + fmt.Println(" /schedule remove - Remove a task by 4-char ID") + fmt.Println(" /schedule clear - Stop and remove all tasks") +} + +func listScheduledTasks() { + tasks := sched.ListTasks() + if len(tasks) == 0 { + fmt.Println("No scheduled tasks") + return + } + + fmt.Printf("\n%-6s %-8s %-8s %-8s %s\n", "ID", "Runs", "Interval", "Last Run", "Output") + fmt.Println(strings.Repeat("-", 70)) + + for _, t := range tasks { + task, ok := t.(*scheduler.ScheduledTask) + if !ok { + continue + } + + interval := fmt.Sprintf("%ds", task.IntervalSec) + if task.RunAt != "" { + interval = task.RunAt + } + + lastRun := "Never" + if !task.LastRun.IsZero() { + lastRun = task.LastRun.Format("15:04") + } + + output := "OK" + if task.LastError != "" { + output = truncate(task.LastError, 40) + } else if task.LastOutput != "" { + lines := strings.Split(task.LastOutput, "\n") + lastLine := strings.TrimSpace(lines[len(lines)-1]) + output = truncate(lastLine, 40) + } + + fmt.Printf("%-6s %-8d %-8s %-8s %s\n", + task.ID, + task.RunCount, + interval, + lastRun, + output, + ) + } + fmt.Println() +} diff --git a/internal/agent/client.go b/internal/agent/client.go index aee8f87..c822c69 100644 --- a/internal/agent/client.go +++ b/internal/agent/client.go @@ -3,7 +3,6 @@ package agent import ( "context" "fmt" - "log" "strings" "sync" @@ -17,6 +16,8 @@ var ( providerCacheMu sync.Mutex mcpExecutors = map[string]*mcp.MCPToolExecutor{} mcpExecutorsMu sync.Mutex + toolsCache = map[string][]providers.Tool{} + toolsCacheMu sync.Mutex ) type SendResult struct { @@ -48,7 +49,6 @@ func SendPrompt(ctx context.Context, provider *config.Provider, prompt string) ( return strings.TrimSpace(resp), nil } - func CancelPrompt(provider *config.Provider) { if provider == nil { return @@ -71,18 +71,22 @@ func CancelPrompt(provider *config.Provider) { mcpExecutorsMu.Unlock() } - func SetProviderTools(provider *config.Provider, tools []providers.Tool) { p, err := GetProvider(provider) if err != nil { return } + if provider != nil { + key := cacheKey(provider) + toolsCacheMu.Lock() + toolsCache[key] = tools + toolsCacheMu.Unlock() + } if pt, ok := p.(interface{ SetTools([]providers.Tool) }); ok { pt.SetTools(tools) } } - func RegisterMCPExecutor(provider *config.Provider, executor *mcp.MCPToolExecutor) { if provider == nil || executor == nil { return @@ -110,7 +114,7 @@ func CloseProvider(cfg *config.Provider) { } providerCacheMu.Unlock() if ok { - log.Printf("[agent] Closing provider for %s/%s", cfg.Type, cfg.Name) + //log.Printf("[agent] Closing provider for %s/%s", cfg.Type, cfg.Name) _ = p.Close() } @@ -140,10 +144,47 @@ func GetProvider(cfg *config.Provider) (providers.Provider, error) { if err != nil { return nil, err } + toolsCacheMu.Lock() + tools := toolsCache[key] + toolsCacheMu.Unlock() + if len(tools) > 0 { + if pt, ok := p.(interface{ SetTools([]providers.Tool) }); ok { + pt.SetTools(tools) + } + } providerCache[key] = p return p, nil } +// ReinitializeProvider closes any cached provider instance for cfg and creates a fresh one. +// This is useful when background workers (like the scheduler) need a clean client instance +// with all settings (API keys, endpoint, params) preserved. +func ReinitializeProvider(cfg *config.Provider) (providers.Provider, error) { + if cfg == nil { + return nil, fmt.Errorf("nil provider") + } + CloseProvider(cfg) + return GetProvider(cfg) +} + func cacheKey(cfg *config.Provider) string { return fmt.Sprintf("%s|%s|%s|%s", cfg.Type, cfg.Name, cfg.Model, cfg.Endpoint) } + +type ProviderConfig struct { + Type string + Name string + Model string +} + +func GetProviderFromConfig(cfg *ProviderConfig) (providers.Provider, error) { + if cfg == nil { + return nil, fmt.Errorf("nil config") + } + provider := &config.Provider{ + Type: cfg.Type, + Name: cfg.Name, + Model: cfg.Model, + } + return GetProvider(provider) +} diff --git a/internal/context/project.go b/internal/context/project.go index fa1267f..c480703 100644 --- a/internal/context/project.go +++ b/internal/context/project.go @@ -11,8 +11,8 @@ import ( ) const ( - maxTreeLines = 500 - treeCacheName = ".gtree" + maxTreeLines = 500 + treeCacheName = ".gtree" ) var ( @@ -81,7 +81,6 @@ func BuildTree(root string) (string, error) { return tree, nil } - func loadGitignore(root string) *ignore.GitIgnore { path := filepath.Join(root, ".gitignore") data, err := os.ReadFile(path) diff --git a/internal/mcp/bash.go b/internal/mcp/bash.go index 2ae681b..98f8a49 100644 --- a/internal/mcp/bash.go +++ b/internal/mcp/bash.go @@ -168,13 +168,13 @@ func (t *JobTracker) JobStatus(id string) (string, bool) { } type BashServer struct { - allowedPaths []string - tools []Tool - jobTracker *JobTracker - autoConfirm bool - mu sync.RWMutex - failedMu sync.Mutex - failedStarts []string + allowedPaths []string + tools []Tool + jobTracker *JobTracker + autoConfirm bool + mu sync.RWMutex + failedMu sync.Mutex + failedStarts []string } var ( @@ -867,7 +867,6 @@ func (s *BashServer) SetAllowedPaths(paths []string) { func (s *BashServer) AllowedPathsUpdated() {} - // Aliases for backwards compatibility with main.go func (s *BashServer) ListBackground() (string, error) { return s.ListBackgroundJobs(), nil diff --git a/internal/mcp/bash_windows.go b/internal/mcp/bash_windows.go index 3bed278..48f7379 100644 --- a/internal/mcp/bash_windows.go +++ b/internal/mcp/bash_windows.go @@ -39,13 +39,11 @@ func setProcessGroup(cmd *exec.Cmd) { cmd.SysProcAttr.CreationFlags |= syscall.CREATE_NEW_PROCESS_GROUP } - // GetUnixTools returns an empty list on Windows (not supported) func GetUnixTools() []Tool { return []Tool{} } - // HandleRunBashScript is not supported on Windows func (s *BashServer) HandleRunBashScript(ctx context.Context, args map[string]interface{}) (string, error) { return "", fmt.Errorf("run_bash_script is not supported on Windows") diff --git a/internal/mcp/scheduler.go b/internal/mcp/scheduler.go new file mode 100644 index 0000000..9ab9c38 --- /dev/null +++ b/internal/mcp/scheduler.go @@ -0,0 +1,207 @@ +package mcp + +import ( + "context" + "encoding/json" + "fmt" + "strings" + "time" +) + +type SchedulerServer struct { + scheduler *schedulerExt +} + +type schedulerExt struct { + scheduler interface { + AddTask(prompt string, intervalSec int, runAt string, providerType, providerName, providerModel string) (interface{}, error) + StopTask(id string) bool + RemoveTask(id string) bool + ListTasks() []interface{} + GetTask(id string) interface{} + } + mu interface { + Lock() + Unlock() + } + provider interface { + GetProvider(cfg interface{}) (interface{}, error) + } + providerMu interface { + Lock() + Unlock() + } +} + +type TaskInfo struct { + ID string `json:"id"` + Prompt string `json:"prompt"` + IntervalSec int `json:"interval_sec"` + RunAt string `json:"run_at"` + IsRepeating bool `json:"is_repeating"` + CreatedAt time.Time `json:"created_at"` + LastRun time.Time `json:"last_run"` + RunCount int `json:"run_count"` + LastError string `json:"last_error"` +} + +func NewSchedulerServer(scheduler interface { + AddTask(prompt string, intervalSec int, runAt string, providerType, providerName, providerModel string) (interface{}, error) + StopTask(id string) bool + RemoveTask(id string) bool + ListTasks() []interface{} + GetTask(id string) interface{} +}, providerGetter interface { + GetProvider(cfg interface{}) (interface{}, error) +}) *SchedulerServer { + return &SchedulerServer{ + scheduler: &schedulerExt{ + scheduler: scheduler, + }, + } +} + +func (s *SchedulerServer) Name() string { + return "scheduler" +} + +func (s *SchedulerServer) Tools() []Tool { + currentTime := time.Now().Format("2006-01-02 15:04:05") + return []Tool{ + { + Name: "scheduler", + Description: fmt.Sprintf("Schedule a task to run a prompt at a specific time or at regular intervals. Current time: %s. Use interval_sec for repeating tasks (e.g., 60 for every minute, 3600 for every hour), or run_at for specific times (e.g., \"14:30\" for 2:30 PM daily, or \"now\" to execute immediately). Set run_once=true to run only once (useful with interval_sec as a one-shot delay).", currentTime), + InputSchema: json.RawMessage(`{ + "type": "object", + "properties": { + "prompt": { + "type": "string", + "description": "The prompt to execute when the task runs" + }, + "interval_sec": { + "type": "integer", + "description": "Interval in seconds for repeating tasks (e.g., 60, 300, 3600). Use 0 or omit for one-time task." + }, + "run_at": { + "type": "string", + "description": "Specific time to run in HH:MM format (e.g., '14:30' for 2:30 PM), or 'now' to execute immediately." + }, + "run_once": { + "type": "boolean", + "description": "If true, the task runs once and does not repeat (even if interval_sec is set)." + } + }, + "required": ["prompt"] + }`), + }, + } +} + +func (s *SchedulerServer) CallTool(ctx context.Context, name string, args map[string]interface{}) (string, error) { + if name != "scheduler" { + return "", fmt.Errorf("unknown tool: %s", name) + } + + prompt, ok := args["prompt"].(string) + if !ok || strings.TrimSpace(prompt) == "" { + return "", fmt.Errorf("prompt is required") + } + + intervalSec := 0 + if iv, ok := args["interval_sec"].(float64); ok { + intervalSec = int(iv) + } else if iv, ok := args["run_every"].(float64); ok { + intervalSec = int(iv) + } + + runAt := "" + if rt, ok := args["run_at"].(string); ok { + runAt = strings.TrimSpace(rt) + } + + runOnce := false + if v, ok := args["run_once"].(bool); ok { + runOnce = v + } else if v, ok := args["run_once"].(string); ok { + runOnce = strings.ToLower(strings.TrimSpace(v)) == "true" + } + + if intervalSec <= 0 && runAt == "" { + return "", fmt.Errorf("must specify either interval_sec or run_at") + } + + isRepeating := intervalSec > 0 + if runOnce { + isRepeating = false + } + if strings.ToLower(runAt) == "now" { + isRepeating = false + } + + var ( + task interface{} + err error + ) + + type schedulerWithRepeat interface { + AddTaskWithRepeat(prompt string, intervalSec int, runAt string, isRepeating bool, providerType, providerName, providerModel string) (interface{}, error) + } + if sr, ok := s.scheduler.scheduler.(schedulerWithRepeat); ok { + task, err = sr.AddTaskWithRepeat(prompt, intervalSec, runAt, isRepeating, "current", "current", "current") + } else { + task, err = s.scheduler.scheduler.AddTask(prompt, intervalSec, runAt, "current", "current", "current") + } + if err != nil { + return fmt.Sprintf("Error scheduling task: %v", err), nil + } + + if task == nil { + return "Error: task was not created", nil + } + + var st struct { + ID string `json:"id"` + Prompt string `json:"prompt"` + IntervalSec int `json:"interval_sec"` + RunAt string `json:"run_at"` + IsRepeating bool `json:"is_repeating"` + } + b, marshalErr := json.Marshal(task) + if marshalErr != nil || json.Unmarshal(b, &st) != nil || strings.TrimSpace(st.ID) == "" { + return "Task scheduled successfully but could not get ID", nil + } + + mode := "repeating" + if !st.IsRepeating { + mode = "once" + } + return fmt.Sprintf("Task scheduled successfully with ID: %s\nMode: %s\nPrompt: %s\nInterval: %d seconds\nRun at: %s", + st.ID, mode, st.Prompt, st.IntervalSec, st.RunAt), nil +} + +func (s *SchedulerServer) AllowedPaths() []string { + return nil +} + +func (s *SchedulerServer) AddPath(ctx context.Context, path string) error { + return fmt.Errorf("scheduler server does not support paths") +} + +func (s *SchedulerServer) TempAddPath(path string) { +} + +func (s *SchedulerServer) AddURL(ctx context.Context, url string) error { + return fmt.Errorf("scheduler server does not support URLs") +} + +func (s *SchedulerServer) RemovePath(ctx context.Context, path string) error { + return fmt.Errorf("scheduler server does not support paths") +} + +func (s *SchedulerServer) RemoveURL(ctx context.Context, url string) error { + return fmt.Errorf("scheduler server does not support URLs") +} + +func (s *SchedulerServer) Close() error { + return nil +} diff --git a/internal/mcp/scheduler_test.go b/internal/mcp/scheduler_test.go new file mode 100644 index 0000000..7146736 --- /dev/null +++ b/internal/mcp/scheduler_test.go @@ -0,0 +1,61 @@ +package mcp + +import ( + "context" + "strings" + "testing" +) + +type fakeScheduler struct { + lastIsRepeating bool +} + +type fakeTask struct { + ID string `json:"id"` + Prompt string `json:"prompt"` + IntervalSec int `json:"interval_sec"` + RunAt string `json:"run_at"` + IsRepeating bool `json:"is_repeating"` +} + +func (f *fakeScheduler) AddTask(prompt string, intervalSec int, runAt string, providerType, providerName, providerModel string) (interface{}, error) { + return f.AddTaskWithRepeat(prompt, intervalSec, runAt, intervalSec > 0, providerType, providerName, providerModel) +} + +func (f *fakeScheduler) AddTaskWithRepeat(prompt string, intervalSec int, runAt string, isRepeating bool, providerType, providerName, providerModel string) (interface{}, error) { + f.lastIsRepeating = isRepeating + return &fakeTask{ID: "ABCD", Prompt: prompt, IntervalSec: intervalSec, RunAt: runAt, IsRepeating: isRepeating}, nil +} + +func (f *fakeScheduler) StopTask(id string) bool { return true } +func (f *fakeScheduler) RemoveTask(id string) bool { return true } +func (f *fakeScheduler) ListTasks() []interface{} { return nil } +func (f *fakeScheduler) GetTask(id string) interface{} { return nil } + +func TestSchedulerTool_RunOnceForcesNonRepeating(t *testing.T) { + fs := &fakeScheduler{} + server := NewSchedulerServer(fs, nil) + + _, err := server.CallTool(context.Background(), "scheduler", map[string]interface{}{ + "prompt": "hi", + "interval_sec": float64(60), + "run_once": true, + }) + if err != nil { + t.Fatalf("CallTool: %v", err) + } + if fs.lastIsRepeating { + t.Fatalf("expected isRepeating=false when run_once=true") + } +} + +func TestSchedulerTool_SchemaIncludesRunOnce(t *testing.T) { + server := NewSchedulerServer(&fakeScheduler{}, nil) + tools := server.Tools() + if len(tools) != 1 { + t.Fatalf("expected 1 tool, got %d", len(tools)) + } + if !strings.Contains(string(tools[0].InputSchema), `"run_once"`) { + t.Fatalf("expected InputSchema to include run_once, got: %s", string(tools[0].InputSchema)) + } +} diff --git a/internal/providers/ollama.go b/internal/providers/ollama.go index e9eed81..681512a 100644 --- a/internal/providers/ollama.go +++ b/internal/providers/ollama.go @@ -220,7 +220,7 @@ func (o *ollamaProvider) Send(ctx context.Context, prompt string) (string, error defer o.sendMu.Unlock() o.mu.Lock() - systemPrompt, userInput := splitSystemUserPrompt(prompt) + systemPrompt, userInput := SplitSystemUserPrompt(prompt) if systemPrompt != "" { o.systemPrompt = systemPrompt } diff --git a/internal/providers/openrouter.go b/internal/providers/openrouter.go index 2a98ce0..c004194 100644 --- a/internal/providers/openrouter.go +++ b/internal/providers/openrouter.go @@ -89,6 +89,19 @@ type openRouterProvider struct { statusCh chan<- string } +func (p *openRouterProvider) emitStatus(msg string) { + p.mu.Lock() + ch := p.statusCh + p.mu.Unlock() + if ch == nil { + return + } + select { + case ch <- msg: + default: + } +} + func init() { Register("openrouter", newOpenRouterProvider) } @@ -162,7 +175,7 @@ func newOpenRouterProvider(cfg *config.Provider) (Provider, error) { } func (p *openRouterProvider) Send(ctx context.Context, prompt string) (string, error) { - systemPrompt, userInput := splitSystemUserPrompt(prompt) + systemPrompt, userInput := SplitSystemUserPrompt(prompt) if userInput == "" { userInput = strings.TrimSpace(prompt) } @@ -219,13 +232,13 @@ func (p *openRouterProvider) Send(ctx context.Context, prompt string) (string, e var lastErr error for attempt := 0; attempt <= maxRetries; attempt++ { if attempt > 0 { - fmt.Printf("\n[OPENROUTER] Waiting %v before retry (attempt %d/%d)...\n", retryDelay, attempt, maxRetries) + p.emitStatus(fmt.Sprintf("OpenRouter: waiting %v before retry (%d/%d)", retryDelay, attempt, maxRetries)) select { case <-ctx.Done(): return "", ctx.Err() case <-time.After(retryDelay): } - fmt.Printf("\n[OPENROUTER] Retrying request (attempt %d/%d)...\n", attempt, maxRetries) + p.emitStatus(fmt.Sprintf("OpenRouter: retrying request (%d/%d)", attempt, maxRetries)) } req, err := http.NewRequestWithContext(ctx, "POST", p.baseURL+"/responses", bytes.NewReader(body)) @@ -242,7 +255,7 @@ func (p *openRouterProvider) Send(ctx context.Context, prompt string) (string, e if err != nil { lastErr = fmt.Errorf("request failed: %w", err) if attempt < maxRetries { - fmt.Printf("[OPENROUTER00] retryable error: %v\n", lastErr) + p.emitStatus(fmt.Sprintf("OpenRouter: retryable error: %v", lastErr)) } continue } @@ -253,7 +266,7 @@ func (p *openRouterProvider) Send(ctx context.Context, prompt string) (string, e if resp.StatusCode == http.StatusTooManyRequests { lastErr = fmt.Errorf("API error (status %d): %s", resp.StatusCode, string(bodyBytes)) if attempt < maxRetries { - fmt.Printf("[OPENROUTER] retryable error: %v\n", lastErr) + p.emitStatus(fmt.Sprintf("OpenRouter: retryable error: %v", lastErr)) } continue } diff --git a/internal/providers/prompt.go b/internal/providers/prompt.go index 706247e..1f3045e 100644 --- a/internal/providers/prompt.go +++ b/internal/providers/prompt.go @@ -7,7 +7,8 @@ const ( userAskedMarker = "User asked:" ) -func splitSystemUserPrompt(prompt string) (string, string) { +// SplitSystemUserPrompt splits a prompt into system and user parts +func SplitSystemUserPrompt(prompt string) (string, string) { prompt = strings.TrimSpace(prompt) if prompt == "" { return "", "" diff --git a/internal/scheduler/scheduler.go b/internal/scheduler/scheduler.go new file mode 100644 index 0000000..94a6c51 --- /dev/null +++ b/internal/scheduler/scheduler.go @@ -0,0 +1,875 @@ +package scheduler + +import ( + "context" + "crypto/rand" + "database/sql" + "encoding/json" + "fmt" + "os" + "path/filepath" + "runtime" + "sort" + "strings" + "sync" + "time" + + "github.com/cheikh2shift/godex/internal/mcp" + "github.com/cheikh2shift/godex/internal/providers" + "github.com/cheikh2shift/godex/internal/toolcalls" + _ "modernc.org/sqlite" +) + +type ScheduledTask struct { + ID string `json:"id"` + Prompt string `json:"prompt"` + IntervalSec int `json:"interval_sec"` + RunAt string `json:"run_at"` + IsRepeating bool `json:"is_repeating"` + WorkingDir string `json:"working_dir"` + CreatedAt time.Time `json:"created_at"` + LastRun time.Time `json:"last_run"` + RunCount int `json:"run_count"` + LastError string `json:"last_error"` + LastOutput string `json:"last_output"` + ProviderType string `json:"provider_type"` + ProviderName string `json:"provider_name"` + ProviderModel string `json:"provider_model"` +} + +type Scheduler struct { + db *sql.DB + tasks map[string]*ScheduledTask + taskMu sync.RWMutex + running map[string]context.CancelFunc + runMu sync.RWMutex + stopCh chan struct{} + wg sync.WaitGroup + provider ProviderGetter + servers []ToolServer + onTaskFinished func(*ScheduledTask) + maxRounds int + toolTimeout int + wd string + os string + statusCh chan string +} + +type ProviderGetter interface { + GetProvider(cfg interface{}) (interface{}, error) +} + +type ToolServer interface { + Name() string + Tools() []mcp.Tool + CallTool(ctx context.Context, name string, args map[string]interface{}) (string, error) +} + +func getDefaultDBPath() string { + if runtime.GOOS == "windows" { + return filepath.Join(os.Getenv("LOCALAPPDATA"), "godex", "scheduler.db") + } + homeDir, err := os.UserHomeDir() + if err != nil || homeDir == "" { + tmpDir := os.TempDir() + return filepath.Join(tmpDir, "godex-scheduler.db") + } + return filepath.Join(homeDir, ".godex", "scheduler.db") +} + +func New(dbPath string) (*Scheduler, error) { + if dbPath == "" { + dbPath = getDefaultDBPath() + } + + wd, _ := os.Getwd() + + dir := filepath.Dir(dbPath) + if err := os.MkdirAll(dir, 0755); err != nil { + return nil, err + } + + db, err := sql.Open("sqlite", dbPath) + if err != nil { + return nil, err + } + + if err := db.Ping(); err != nil { + return nil, err + } + + s := &Scheduler{ + db: db, + tasks: make(map[string]*ScheduledTask), + running: make(map[string]context.CancelFunc), + stopCh: make(chan struct{}), + wd: wd, + statusCh: make(chan string, 8), + } + + if err := s.initSchema(); err != nil { + db.Close() + return nil, err + } + + if err := s.loadTasks(); err != nil { + db.Close() + return nil, err + } + + s.startLoadedTasks() + + return s, nil +} + +func (s *Scheduler) startLoadedTasks() { + s.taskMu.Lock() + defer s.taskMu.Unlock() + for _, task := range s.tasks { + if task.IsRepeating || (task.LastRun.IsZero() && task.IntervalSec > 0) { + s.startTask(task) + } + } +} + +func NewDefault() (*Scheduler, error) { + return New("") +} + +func (s *Scheduler) SetStatusChannel(ch chan string) { + s.statusCh = ch +} + +func (s *Scheduler) initSchema() error { + schema := ` + CREATE TABLE IF NOT EXISTS scheduled_tasks ( + id TEXT PRIMARY KEY, + prompt TEXT NOT NULL, + interval_sec INTEGER NOT NULL, + run_at TEXT NOT NULL, + is_repeating INTEGER NOT NULL, + working_dir TEXT NOT NULL, + created_at TEXT NOT NULL, + last_run TEXT, + run_count INTEGER DEFAULT 0, + last_error TEXT, + last_output TEXT, + provider_type TEXT NOT NULL, + provider_name TEXT NOT NULL, + provider_model TEXT NOT NULL + ); + CREATE INDEX IF NOT EXISTS idx_scheduled_tasks_created_at ON scheduled_tasks(created_at); + CREATE INDEX IF NOT EXISTS idx_scheduled_tasks_working_dir ON scheduled_tasks(working_dir); + ` + _, err := s.db.Exec(schema) + if err != nil { + return err + } + + // Migrate older DBs which predate working_dir. + if _, err := s.db.Exec(`ALTER TABLE scheduled_tasks ADD COLUMN working_dir TEXT`); err != nil { + // ignore "duplicate column name" errors (SQLite reports this when column already exists) + if !strings.Contains(strings.ToLower(err.Error()), "duplicate column") { + return err + } + } + _, _ = s.db.Exec(`CREATE INDEX IF NOT EXISTS idx_scheduled_tasks_working_dir ON scheduled_tasks(working_dir)`) + if strings.TrimSpace(s.wd) != "" { + _, _ = s.db.Exec(`UPDATE scheduled_tasks SET working_dir = ? WHERE working_dir IS NULL OR working_dir = ''`, s.wd) + } + + return nil +} + +func (s *Scheduler) loadTasks() error { + query := ` + SELECT id, prompt, interval_sec, run_at, is_repeating, working_dir, created_at, + last_run, run_count, last_error, last_output, provider_type, provider_name, provider_model + FROM scheduled_tasks + ` + if strings.TrimSpace(s.wd) != "" { + query += ` WHERE working_dir = ?` + } + query += ` ORDER BY created_at ASC` + + var ( + rows *sql.Rows + err error + ) + if strings.TrimSpace(s.wd) != "" { + rows, err = s.db.Query(query, s.wd) + } else { + rows, err = s.db.Query(query) + } + if err != nil { + return err + } + defer rows.Close() + + s.taskMu.Lock() + defer s.taskMu.Unlock() + + for rows.Next() { + var task ScheduledTask + var ( + lastRun, createdAt, runAt, lastOutput sql.NullString + workingDir sql.NullString + isRepeatingInt sql.NullInt64 + ) + err := rows.Scan( + &task.ID, &task.Prompt, &task.IntervalSec, &runAt, &isRepeatingInt, &workingDir, + &createdAt, &lastRun, &task.RunCount, &task.LastError, &lastOutput, + &task.ProviderType, &task.ProviderName, &task.ProviderModel, + ) + if err != nil { + return err + } + if isRepeatingInt.Valid && isRepeatingInt.Int64 != 0 { + task.IsRepeating = true + } + if createdAt.Valid { + task.CreatedAt, _ = time.Parse(time.RFC3339, createdAt.String) + } + if lastRun.Valid { + task.LastRun, _ = time.Parse(time.RFC3339, lastRun.String) + } + if runAt.Valid { + task.RunAt = runAt.String + } + if lastOutput.Valid { + task.LastOutput = lastOutput.String + } + if workingDir.Valid { + task.WorkingDir = workingDir.String + } else { + task.WorkingDir = s.wd + } + s.tasks[task.ID] = &task + } + return rows.Err() +} + +func (s *Scheduler) SetProviderGetter(pg ProviderGetter) { + s.provider = pg +} + +func (s *Scheduler) SetServers(servers []interface{}) { + s.servers = nil + for _, srv := range servers { + ts, ok := srv.(ToolServer) + if !ok { + continue + } + s.servers = append(s.servers, ts) + } +} + +func (s *Scheduler) SetOnTaskFinished(cb func(*ScheduledTask)) { + s.onTaskFinished = cb +} + +func (s *Scheduler) SetConfig(maxRounds, toolTimeout int, wd, os string) { + if maxRounds > 0 { + s.maxRounds = maxRounds + } + if toolTimeout > 0 { + s.toolTimeout = toolTimeout + } + if strings.TrimSpace(wd) != "" && s.wd != wd { + // Re-scope tasks to the configured working directory. + s.wd = wd + s.StopAllTasks() + s.taskMu.Lock() + s.tasks = make(map[string]*ScheduledTask) + s.taskMu.Unlock() + _ = s.initSchema() + _ = s.loadTasks() + s.startLoadedTasks() + } else if strings.TrimSpace(wd) != "" { + s.wd = wd + } + if os != "" { + s.os = os + } +} + +func (s *Scheduler) generateID() (string, error) { + charset := "ABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789" + bytes := make([]byte, 4) + if _, err := rand.Read(bytes); err != nil { + return "", err + } + result := make([]byte, 4) + for i, b := range bytes { + result[i] = charset[int(b)%len(charset)] + } + return string(result), nil +} + +type SchedulerInterface interface { + AddTask(prompt string, intervalSec int, runAt string, providerType, providerName, providerModel string) (interface{}, error) + AddTaskWithRepeat(prompt string, intervalSec int, runAt string, isRepeating bool, providerType, providerName, providerModel string) (interface{}, error) + StopTask(id string) bool + RemoveTask(id string) bool + ListTasks() []interface{} + GetTask(id string) interface{} +} + +func (s *Scheduler) AddTask(prompt string, intervalSec int, runAt string, providerType, providerName, providerModel string) (interface{}, error) { + isRepeating := intervalSec > 0 + return s.AddTaskWithRepeat(prompt, intervalSec, runAt, isRepeating, providerType, providerName, providerModel) +} + +func (s *Scheduler) AddTaskWithRepeat(prompt string, intervalSec int, runAt string, isRepeating bool, providerType, providerName, providerModel string) (interface{}, error) { + if intervalSec <= 0 && runAt == "" { + return nil, fmt.Errorf("must specify either interval or run_at time") + } + + id, err := s.generateID() + if err != nil { + return nil, err + } + + task := &ScheduledTask{ + ID: id, + Prompt: prompt, + IntervalSec: intervalSec, + RunAt: runAt, + IsRepeating: isRepeating, + WorkingDir: s.wd, + CreatedAt: time.Now(), + RunCount: 0, + ProviderType: providerType, + ProviderName: providerName, + ProviderModel: providerModel, + } + + s.taskMu.Lock() + s.tasks[id] = task + s.taskMu.Unlock() + + if err := s.saveTask(task); err != nil { + return nil, err + } + + s.startTask(task) + + return task, nil +} + +func (s *Scheduler) saveTask(task *ScheduledTask) error { + var lastRunStr string + if task.LastRun.IsZero() { + lastRunStr = "" + } else { + lastRunStr = task.LastRun.Format(time.RFC3339) + } + _, err := s.db.Exec(` + INSERT OR REPLACE INTO scheduled_tasks + (id, prompt, interval_sec, run_at, is_repeating, working_dir, created_at, last_run, run_count, last_error, last_output, provider_type, provider_name, provider_model) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + `, task.ID, task.Prompt, task.IntervalSec, task.RunAt, task.IsRepeating, task.WorkingDir, + task.CreatedAt.Format(time.RFC3339), lastRunStr, + task.RunCount, task.LastError, task.LastOutput, task.ProviderType, task.ProviderName, task.ProviderModel) + return err +} + +func (s *Scheduler) startTask(task *ScheduledTask) { + s.runMu.Lock() + if _, exists := s.running[task.ID]; exists { + s.runMu.Unlock() + return + } + + ctx, cancel := context.WithCancel(context.Background()) + s.running[task.ID] = cancel + s.runMu.Unlock() + + s.wg.Add(1) + go s.runLoop(ctx, task) +} + +func (s *Scheduler) runLoop(ctx context.Context, task *ScheduledTask) { + defer s.wg.Done() + + sleepDuration := func() time.Duration { + if strings.ToLower(strings.TrimSpace(task.RunAt)) == "now" { + return 0 + } + if task.IntervalSec > 0 { + return time.Duration(task.IntervalSec) * time.Second + } + if task.RunAt != "" { + now := time.Now() + target, err := time.Parse("15:04", task.RunAt) + if err != nil { + return 24 * time.Hour + } + target = time.Date(now.Year(), now.Month(), now.Day(), target.Hour(), target.Minute(), 0, 0, now.Location()) + if target.Before(now) { + target = target.Add(24 * time.Hour) + } + return target.Sub(now) + } + return 24 * time.Hour + } + + for { + select { + case <-ctx.Done(): + return + case <-time.After(sleepDuration()): + s.executeTask(task) + if !task.IsRepeating { + return + } + } + } +} + +func (s *Scheduler) executeTask(task *ScheduledTask) { + taskCtx, cancel := context.WithTimeout(context.Background(), 5*time.Minute) + defer cancel() + + if len(s.servers) > 0 { + result, err := s.executeWithTools(taskCtx, task) + if err != nil { + task.LastError = fmt.Sprintf("execution error: %v", err) + task.LastOutput = "" + } else { + task.LastError = "" + task.LastOutput = result + } + } else if s.provider != nil { + var cfg interface{} + if task.ProviderType == "unknown" || task.ProviderType == "" || task.ProviderType == "current" { + cfg = map[string]interface{}{ + "type": "current", + "name": "current", + "model": "current", + } + } else { + cfg = map[string]interface{}{ + "type": task.ProviderType, + "name": task.ProviderName, + "model": task.ProviderModel, + } + } + + prov, err := s.provider.GetProvider(cfg) + if err != nil { + task.LastError = fmt.Sprintf("failed to get provider: %v", err) + task.LastRun = time.Now() + task.RunCount++ + s.updateTask(task) + return + } + + provInterface, ok := prov.(interface { + Send(ctx context.Context, prompt string) (string, error) + }) + if !ok { + task.LastError = "provider does not support Send method" + task.LastRun = time.Now() + task.RunCount++ + s.updateTask(task) + return + } + + result, err := provInterface.Send(taskCtx, task.Prompt) + if err != nil { + task.LastError = fmt.Sprintf("execution error: %v", err) + task.LastOutput = "" + } else { + task.LastError = "" + task.LastOutput = result + } + } else { + task.LastError = "no provider or servers available" + } + + task.LastRun = time.Now() + task.RunCount++ + s.updateTask(task) + + if s.onTaskFinished != nil { + func() { + defer func() { _ = recover() }() + s.onTaskFinished(task) + }() + } +} + +func (s *Scheduler) executeWithTools(ctx context.Context, task *ScheduledTask) (string, error) { + maxRounds := s.maxRounds + if maxRounds <= 0 { + maxRounds = 10 + } + toolTimeout := s.toolTimeout + if toolTimeout <= 0 { + toolTimeout = 180 + } + + if s.provider == nil { + return "", fmt.Errorf("no provider available") + } + + cfg := map[string]interface{}{ + "type": "current", + "name": "current", + "model": "current", + } + prov, err := s.provider.GetProvider(cfg) + if err != nil { + return "", fmt.Errorf("failed to get provider: %v", err) + } + llmProvider, ok := prov.(interface { + Send(ctx context.Context, prompt string) (string, error) + SetThinkCallback(cb func(string)) + SupportsNativeToolCalls() bool + }) + if !ok { + return "", fmt.Errorf("provider does not support required interface") + } + + toolsDesc := s.getToolsDescription() + toolCallFormat := "" + if !llmProvider.SupportsNativeToolCalls() { + toolCallFormat = "\nTo call tools, respond with JSON in markdown code blocks:\n```json\n{\n \"name\": \"tool_name\",\n \"arguments\": {\n \"arg1\": \"value1\"\n }\n}\n```" + } + + input := task.Prompt + fullPrompt := s.buildInitialPrompt(input, toolsDesc, toolCallFormat, llmProvider.SupportsNativeToolCalls()) + + prevRoundToolCalls := make(map[string]bool) + prevNoTool := false + + for round := 0; round < maxRounds; round++ { + var resp string + resp, err = llmProvider.Send(ctx, fullPrompt) + if err != nil { + return "", fmt.Errorf("provider error: %v", err) + } + + preResp, finalResp, hasFinal := s.splitFinalAnswer(resp) + toolCalls := toolcalls.ExtractAllToolCalls(preResp) + isToolCallResponse := len(toolCalls) > 0 + + if !isToolCallResponse || len(toolCalls) == 0 { + if !hasFinal { + if !s.looksLikeToolCall(resp) && round > 0 { + return resp, nil + } + if strings.TrimSpace(resp) != "" && prevNoTool { + return resp, nil + } + prevNoTool = true + continuePrompt := fmt.Sprintf("You provided the following response:\n%s\n\nHowever, I couldn't find any valid tool calls in it. Please provide your FINAL_ANSWER now based on the information you have.", resp) + fullPrompt = s.buildContinuePrompt(continuePrompt, toolsDesc, toolCallFormat, llmProvider.SupportsNativeToolCalls(), round+1, maxRounds) + continue + } + return finalResp, nil + } + + prevNoTool = false + + var toolResults []string + seenThisRound := make(map[string]bool) + for _, tc := range toolCalls { + toolName, ok := tc["name"].(string) + if !ok { + continue + } + args, ok := tc["arguments"].(map[string]interface{}) + if !ok { + continue + } + argsJSON, _ := json.Marshal(args) + sig := fmt.Sprintf("%s:%s", toolName, string(argsJSON)) + if seenThisRound[sig] { + continue + } + seenThisRound[sig] = true + + if s.statusCh != nil { + select { + case s.statusCh <- fmt.Sprintf("[%s] Calling tool: %s", task.ID, toolName): + //fmt.Printf("DEBUG: Sent status: [%s] Calling tool: %s\n", task.ID, toolName) + default: + } + } + + result, err := s.callTool(ctx, toolName, args, toolTimeout) + if err != nil { + errStr := strings.ToLower(err.Error()) + if strings.Contains(errStr, "path_restricted") || + strings.Contains(errStr, "not allowed") || + strings.Contains(errStr, "denied") || + strings.Contains(errStr, "_blocked") || + strings.Contains(errStr, "permission") { + return fmt.Sprintf("PERMISSION_DENIED: %v", err), nil + } + toolResults = append(toolResults, fmt.Sprintf("ERROR: %v", err)) + } else { + toolResults = append(toolResults, s.truncate(result, 2500)) + } + } + + if hasFinal { + return finalResp, nil + } + + currentToolCalls := make(map[string]bool) + for _, tc := range toolCalls { + name, ok := tc["name"].(string) + if !ok { + continue + } + args := tc["arguments"] + argsJSON, _ := json.Marshal(args) + currentToolCalls[fmt.Sprintf("%s:%s", name, string(argsJSON))] = true + } + + hasRepeatedCalls := false + for sig := range currentToolCalls { + if prevRoundToolCalls[sig] { + hasRepeatedCalls = true + break + } + } + + duplicateWarning := "" + if hasRepeatedCalls && round > 0 { + duplicateWarning = "\n\nWARNING: You are calling the same tool(s) with the same arguments as the previous round. Do NOT call them again." + } + + roundUrgency := "" + if round >= 10 { + roundUrgency = "\n\nSTOP: You have reached round " + fmt.Sprintf("%d", round+1) + ". You MUST provide your FINAL_ANSWER now." + } else if round >= 2 { + roundUrgency = "\n\nNOTE: You are on round " + fmt.Sprintf("%d", round+1) + ". Only call more tools if absolutely necessary." + } + + prevRoundToolCalls = currentToolCalls + + if llmProvider.SupportsNativeToolCalls() { + fullPrompt = fmt.Sprintf("User asked: %s\n\nTool results:\n%s%s%s\n\nProvide your FINAL_ANSWER now.", input, strings.Join(toolResults, "\n---\n"), duplicateWarning, roundUrgency) + } else { + fullPrompt = fmt.Sprintf("You have access to these tools:%s\n\n%s\n\nUser asked: %s\n\nTool results:\n%s%s%s\n\nProvide your FINAL_ANSWER now.", toolsDesc, toolCallFormat, input, strings.Join(toolResults, "\n---\n"), duplicateWarning, roundUrgency) + } + } + + return "Max tool rounds reached", nil +} + +func (s *Scheduler) buildInitialPrompt(input, toolsDesc string, toolCallFormat string, nativeToolCalls bool) string { + osInfo := "linux" + if s.os != "" { + osInfo = s.os + } + arch := runtime.GOARCH + wdInfo := "" + if s.wd != "" { + wdInfo = s.wd + } + + base := fmt.Sprintf(`CRITICAL INFORMATION: +- Operating System: %s (%s) +- Current working directory: %s +Use this path when the user asks about "this folder", "current directory", or similar. + +IMPORTANT: Execute tools FIRST, perform any action asked for by the user, then provide the final answer. Do NOT include any final answer, summary, or "FINAL_ANSWER:" until AFTER you have executed all necessary tools and received their results. +`, osInfo, arch, wdInfo) + + if nativeToolCalls { + return base + "\n\n" + input + } + return fmt.Sprintf("%s\n\nYou have access to these tools:\n%s\n\n%s\n\nUser asked: %s", base, toolsDesc, toolCallFormat, input) +} + +func (s *Scheduler) buildContinuePrompt(continuePrompt, toolsDesc string, toolCallFormat string, nativeToolCalls bool, currentRound, maxRounds int) string { + roundInfo := "" + if currentRound > 0 { + roundInfo = fmt.Sprintf("\n\nNOTE: You are on round %d/%d. Only call more tools if absolutely necessary.", currentRound, maxRounds) + } + osInfo := "linux" + if s.os != "" { + osInfo = s.os + } + arch := runtime.GOARCH + wdInfo := "" + if s.wd != "" { + wdInfo = s.wd + } + extraInfo := fmt.Sprintf(`CRITICAL INFORMATION: +- Operating System: %s (%s) +- Current working directory: %s +Use this path when the user asks about "this folder", "current directory", or similar. +`, osInfo, arch, wdInfo) + + systemPrompt, _ := providers.SplitSystemUserPrompt(continuePrompt) + if systemPrompt != "" { + systemPrompt += "\n\n" + } + userPromptWithContext := fmt.Sprintf("\n\n%s%sContinue from where you left off. %s\n\nProvide your FINAL_ANSWER now.%s", extraInfo, systemPrompt, continuePrompt, roundInfo) + + if nativeToolCalls { + return userPromptWithContext + } + + return fmt.Sprintf("%s\n\nYou have access to these tools:\n%s\n\n%s\n\n%s", extraInfo, toolsDesc, toolCallFormat, userPromptWithContext) +} + +func (s *Scheduler) splitFinalAnswer(text string) (string, string, bool) { + parts := strings.Split(text, "FINAL_ANSWER:") + if len(parts) > 1 { + return parts[0], strings.TrimSpace(parts[1]), true + } + parts = strings.Split(text, "final answer:") + if len(parts) > 1 { + return parts[0], strings.TrimSpace(parts[1]), true + } + return text, "", false +} + +func (s *Scheduler) looksLikeToolCall(text string) bool { + return strings.Contains(text, "tool") || strings.Contains(text, "call") || + strings.Contains(text, "function") || strings.Contains(text, "(") +} + +func (s *Scheduler) callTool(ctx context.Context, name string, args map[string]interface{}, timeoutSecs int) (string, error) { + for _, server := range s.servers { + for _, t := range server.Tools() { + if t.Name == name { + toolCtx, cancel := context.WithTimeout(ctx, time.Duration(timeoutSecs)*time.Second) + defer cancel() + return server.CallTool(toolCtx, name, args) + } + } + } + + // Compatibility: some models incorrectly use the server name (e.g. "bash") as the tool name. + lower := strings.ToLower(strings.TrimSpace(name)) + for _, server := range s.servers { + if strings.ToLower(server.Name()) != lower { + continue + } + for _, t := range server.Tools() { + if t.Name == "run_command" { + toolCtx, cancel := context.WithTimeout(ctx, time.Duration(timeoutSecs)*time.Second) + defer cancel() + return server.CallTool(toolCtx, t.Name, args) + } + } + } + + return "", fmt.Errorf("tool %s not found", name) +} + +func (s *Scheduler) truncate(str string, maxLen int) string { + if len(str) <= maxLen { + return str + } + return str[:maxLen] + "..." +} + +func (s *Scheduler) getToolsDescription() string { + if len(s.servers) == 0 { + return "No tools available." + } + + var desc strings.Builder + for _, server := range s.servers { + tools := server.Tools() + if len(tools) == 0 { + continue + } + desc.WriteString(fmt.Sprintf("\n[%s]\n", server.Name())) + for _, tool := range tools { + desc.WriteString(fmt.Sprintf(" - %s: %s\n", tool.Name, tool.Description)) + } + } + return desc.String() +} + +func (s *Scheduler) updateTask(task *ScheduledTask) { + s.taskMu.Lock() + s.tasks[task.ID] = task + s.taskMu.Unlock() + s.saveTask(task) +} + +func (s *Scheduler) StopTask(id string) bool { + s.runMu.Lock() + if cancel, exists := s.running[id]; exists { + cancel() + delete(s.running, id) + s.runMu.Unlock() + return true + } + s.runMu.Unlock() + return false +} + +func (s *Scheduler) StopAllTasks() { + s.runMu.Lock() + for id, cancel := range s.running { + cancel() + delete(s.running, id) + } + s.runMu.Unlock() +} + +func (s *Scheduler) RemoveTask(id string) bool { + s.StopTask(id) + s.taskMu.Lock() + if _, exists := s.tasks[id]; exists { + delete(s.tasks, id) + s.taskMu.Unlock() + if strings.TrimSpace(s.wd) != "" { + _, err := s.db.Exec("DELETE FROM scheduled_tasks WHERE id = ? AND working_dir = ?", id, s.wd) + return err == nil + } + _, err := s.db.Exec("DELETE FROM scheduled_tasks WHERE id = ?", id) + return err == nil + } + s.taskMu.Unlock() + return false +} + +func (s *Scheduler) ListTasks() []interface{} { + s.taskMu.RLock() + defer s.taskMu.RUnlock() + + tasks := make([]*ScheduledTask, 0, len(s.tasks)) + for _, task := range s.tasks { + tasks = append(tasks, task) + } + sort.Slice(tasks, func(i, j int) bool { + return tasks[i].CreatedAt.Before(tasks[j].CreatedAt) + }) + result := make([]interface{}, len(tasks)) + for i, t := range tasks { + result[i] = t + } + return result +} + +func (s *Scheduler) GetTask(id string) interface{} { + s.taskMu.RLock() + defer s.taskMu.RUnlock() + return s.tasks[id] +} + +func (s *Scheduler) Close() { + close(s.stopCh) + s.StopAllTasks() + s.wg.Wait() + s.db.Close() +} + +type providerConfig struct { + Type string `json:"type"` + Name string `json:"name"` + Model string `json:"model"` +} diff --git a/internal/scheduler/scheduler_callback_test.go b/internal/scheduler/scheduler_callback_test.go new file mode 100644 index 0000000..6dfffc2 --- /dev/null +++ b/internal/scheduler/scheduler_callback_test.go @@ -0,0 +1,40 @@ +package scheduler + +import ( + "sync/atomic" + "testing" + "time" +) + +func TestScheduler_OnTaskFinishedCalled(t *testing.T) { + dbPath := t.TempDir() + "/scheduler.db" + s, err := New(dbPath) + if err != nil { + t.Fatalf("New: %v", err) + } + defer s.Close() + + var called int32 + s.SetOnTaskFinished(func(*ScheduledTask) { + atomic.AddInt32(&called, 1) + }) + + task := &ScheduledTask{ + ID: "TEST", + Prompt: "noop", + IntervalSec: 0, + RunAt: "", + IsRepeating: false, + WorkingDir: s.wd, + CreatedAt: time.Now(), + ProviderType: "current", + ProviderName: "current", + ProviderModel: "current", + } + + s.executeTask(task) + + if atomic.LoadInt32(&called) != 1 { + t.Fatalf("expected callback called once, got %d", atomic.LoadInt32(&called)) + } +} diff --git a/internal/scheduler/scheduler_dedupe_test.go b/internal/scheduler/scheduler_dedupe_test.go new file mode 100644 index 0000000..d16ed4d --- /dev/null +++ b/internal/scheduler/scheduler_dedupe_test.go @@ -0,0 +1,80 @@ +package scheduler + +import ( + "context" + "testing" + "time" + + "github.com/cheikh2shift/godex/internal/mcp" +) + +type fakeProviderGetter struct { + p interface{} +} + +func (f *fakeProviderGetter) GetProvider(cfg interface{}) (interface{}, error) { + return f.p, nil +} + +type fakeLLMProvider struct { + resp string +} + +func (f *fakeLLMProvider) Send(ctx context.Context, prompt string) (string, error) { + return f.resp, nil +} +func (f *fakeLLMProvider) SetThinkCallback(cb func(string)) {} +func (f *fakeLLMProvider) SupportsNativeToolCalls() bool { return false } + +type countingToolServer struct { + name string + tools []mcp.Tool + calls int +} + +func (c *countingToolServer) Name() string { return c.name } +func (c *countingToolServer) Tools() []mcp.Tool { return c.tools } +func (c *countingToolServer) CallTool(ctx context.Context, name string, args map[string]interface{}) (string, error) { + c.calls++ + return "ok", nil +} + +func TestExecuteWithTools_DedupesIdenticalToolCallsInOneResponse(t *testing.T) { + s := &Scheduler{} + server := &countingToolServer{ + name: "scheduler", + tools: []mcp.Tool{ + {Name: "scheduler", Description: "schedule"}, + }, + } + s.servers = []ToolServer{server} + s.provider = &fakeProviderGetter{p: &fakeLLMProvider{ + // Two identical tool calls in one response. Should execute only once. + resp: "```json\n" + + "{\"name\":\"scheduler\",\"arguments\":{\"prompt\":\"x\",\"interval_sec\":60}}\n" + + "```\n" + + "```json\n" + + "{\"name\":\"scheduler\",\"arguments\":{\"prompt\":\"x\",\"interval_sec\":60}}\n" + + "```", + }} + s.maxRounds = 1 + s.toolTimeout = 1 + + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + + task := &ScheduledTask{ + Prompt: "ignored", + ProviderType: "current", + ProviderName: "current", + ProviderModel: "current", + } + + _, err := s.executeWithTools(ctx, task) + if err != nil { + t.Fatalf("executeWithTools: %v", err) + } + if server.calls != 1 { + t.Fatalf("expected 1 tool execution, got %d", server.calls) + } +} diff --git a/internal/scheduler/scheduler_test.go b/internal/scheduler/scheduler_test.go new file mode 100644 index 0000000..81ce873 --- /dev/null +++ b/internal/scheduler/scheduler_test.go @@ -0,0 +1,95 @@ +package scheduler + +import ( + "context" + "strings" + "testing" + "time" + + "github.com/cheikh2shift/godex/internal/mcp" +) + +type fakeToolServer struct { + name string + tools []mcp.Tool +} + +func (f *fakeToolServer) Name() string { return f.name } +func (f *fakeToolServer) Tools() []mcp.Tool { + return f.tools +} +func (f *fakeToolServer) CallTool(ctx context.Context, name string, args map[string]interface{}) (string, error) { + if name == "run_command" { + return "ok", nil + } + return "bad:" + name, nil +} + +func TestScheduler_GetToolsDescription_MatchesMainFormat(t *testing.T) { + s := &Scheduler{} + s.SetServers([]interface{}{ + &fakeToolServer{ + name: "bash", + tools: []mcp.Tool{ + {Name: "run_command", Description: "Run a shell command"}, + }, + }, + }) + + desc := s.getToolsDescription() + if desc == "" { + t.Fatalf("expected non-empty tools description") + } + if want := "[bash]"; !strings.Contains(desc, want) { + t.Fatalf("expected tools description to contain %q, got: %s", want, desc) + } + if want := "run_command"; !strings.Contains(desc, want) { + t.Fatalf("expected tools description to contain %q, got: %s", want, desc) + } +} + +func TestScheduler_CallTool_FindsTool(t *testing.T) { + s := &Scheduler{} + s.SetServers([]interface{}{ + &fakeToolServer{ + name: "bash", + tools: []mcp.Tool{ + {Name: "run_command", Description: "Run a shell command"}, + }, + }, + }) + + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + + out, err := s.callTool(ctx, "run_command", map[string]interface{}{"command": "echo hi"}, 1) + if err != nil { + t.Fatalf("expected nil error, got %v", err) + } + if out != "ok" { + t.Fatalf("expected ok, got %q", out) + } +} + +func TestScheduler_CallTool_ServerNameCompatibility(t *testing.T) { + s := &Scheduler{} + s.SetServers([]interface{}{ + &fakeToolServer{ + name: "bash", + tools: []mcp.Tool{ + {Name: "run_command", Description: "Run a shell command"}, + }, + }, + }) + + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + + out, err := s.callTool(ctx, "Bash", map[string]interface{}{"command": "echo hi"}, 1) + if err != nil { + t.Fatalf("expected nil error, got %v", err) + } + if out != "ok" { + t.Fatalf("expected ok, got %q", out) + } +} diff --git a/internal/scheduler/scheduler_wd_test.go b/internal/scheduler/scheduler_wd_test.go new file mode 100644 index 0000000..9db89d7 --- /dev/null +++ b/internal/scheduler/scheduler_wd_test.go @@ -0,0 +1,73 @@ +package scheduler + +import ( + "os" + "path/filepath" + "testing" +) + +func TestScheduler_ScopesTasksToWorkingDir(t *testing.T) { + tmp := t.TempDir() + dbPath := filepath.Join(tmp, "scheduler.db") + + wd1 := filepath.Join(tmp, "repo1") + wd2 := filepath.Join(tmp, "repo2") + + // Create a repeating task scoped to wd1. + s1, err := New(dbPath) + if err != nil { + t.Fatalf("New: %v", err) + } + s1.SetConfig(0, 0, wd1, "") + if _, err := s1.AddTask("hello", 60, "", "current", "current", "current"); err != nil { + t.Fatalf("AddTask: %v", err) + } + s1.Close() + + // Re-open with the same wd; task should be visible and repeating. + s2, err := New(dbPath) + if err != nil { + t.Fatalf("New (2): %v", err) + } + s2.SetConfig(0, 0, wd1, "") + tasks := s2.ListTasks() + if len(tasks) != 1 { + t.Fatalf("expected 1 task for wd1, got %d", len(tasks)) + } + task, ok := tasks[0].(*ScheduledTask) + if !ok { + t.Fatalf("expected *ScheduledTask, got %T", tasks[0]) + } + if task.WorkingDir != wd1 { + t.Fatalf("expected WorkingDir=%q, got %q", wd1, task.WorkingDir) + } + if !task.IsRepeating { + t.Fatalf("expected IsRepeating=true after reload") + } + s2.Close() + + // Re-open with a different wd; task should not be visible. + s3, err := New(dbPath) + if err != nil { + t.Fatalf("New (3): %v", err) + } + s3.SetConfig(0, 0, wd2, "") + if got := len(s3.ListTasks()); got != 0 { + t.Fatalf("expected 0 tasks for wd2, got %d", got) + } + s3.Close() +} + +func TestScheduler_DefaultWDUsesProcessWD(t *testing.T) { + dbPath := filepath.Join(t.TempDir(), "scheduler.db") + s, err := New(dbPath) + if err != nil { + t.Fatalf("New: %v", err) + } + defer s.Close() + + pwd, _ := os.Getwd() + if pwd != "" && s.wd != pwd { + t.Fatalf("expected scheduler wd to default to process wd %q, got %q", pwd, s.wd) + } +} diff --git a/internal/toolcalls/toolcalls.go b/internal/toolcalls/toolcalls.go new file mode 100644 index 0000000..9471e08 --- /dev/null +++ b/internal/toolcalls/toolcalls.go @@ -0,0 +1,369 @@ +package toolcalls + +import ( + "encoding/json" + "regexp" + "strings" +) + +// ExtractAllToolCalls extracts tool calls from a model response in a variety of common formats. +// It returns items shaped as: {"name": string, "arguments": map[string]interface{}}. +func ExtractAllToolCalls(text string) []map[string]interface{} { + var results []map[string]interface{} + + candidates := []string{text} + if normalized := normalizeToolCallText(text); normalized != text { + candidates = append(candidates, normalized) + } + + for _, candidate := range candidates { + // First, try to find markdown code blocks using proper extraction + codeBlocks := extractCodeBlockContent(candidate) + + for _, content := range codeBlocks { + content = strings.TrimSpace(content) + + // Try to parse the whole block as one JSON object + var data map[string]interface{} + if err := json.Unmarshal([]byte(content), &data); err == nil { + if _, ok := data["name"].(string); ok { + results = append(results, processToolData(data)) + continue + } + } + + // Try to parse the block as a JSON array of tool calls + if arrayCalls := extractToolCallsFromJSONArray(content); len(arrayCalls) > 0 { + results = append(results, arrayCalls...) + continue + } + + // If not a single object, try sanitizing multi-line strings + sanitized := sanitizeMultilineJson(content) + var sanitizedData map[string]interface{} + if err := json.Unmarshal([]byte(sanitized), &sanitizedData); err == nil { + if _, ok := sanitizedData["name"].(string); ok { + results = append(results, processToolData(sanitizedData)) + continue + } + } + if arrayCalls := extractToolCallsFromJSONArray(sanitized); len(arrayCalls) > 0 { + results = append(results, arrayCalls...) + continue + } + + // Fall back to extracting individual JSON objects + extracted := extractJsonObjects(sanitized) + if len(extracted) > 0 { + results = append(results, extracted...) + } + } + + // If no tools found in code blocks, search the whole text for { ... } blobs + if len(results) == 0 { + results = append(results, extractJsonObjects(candidate)...) + } + if len(results) == 0 { + results = append(results, extractToolCallsFromJSONArray(candidate)...) + } + + // Safety fallback for the specific [TOOL_CALL] format + if len(results) == 0 { + toolCallRe := regexp.MustCompile(`(?s)\[TOOL_CALL\]\s*(.*?)\s*\[/TOOL_CALL\]`) + tcMatches := toolCallRe.FindAllStringSubmatch(candidate, -1) + for _, m := range tcMatches { + if len(m) > 1 { + nameRe := regexp.MustCompile(`tool\s*=>\s*"([^"]+)"`) + if nameMatch := nameRe.FindStringSubmatch(m[1]); len(nameMatch) > 1 { + results = append(results, map[string]interface{}{"name": nameMatch[1], "arguments": map[string]interface{}{}}) + } + } + } + } + + // Parse native tool call format: [TOOL_CALL: name | {"arg": "value"}] + if len(results) == 0 { + nativeMatches := extractNativeToolCalls(candidate) + for _, m := range nativeMatches { + name := m[0] + argsStr := m[1] + var args map[string]interface{} + if err := json.Unmarshal([]byte(argsStr), &args); err != nil { + args = map[string]interface{}{"_raw": argsStr} + } + results = append(results, map[string]interface{}{"name": name, "arguments": args}) + } + } + + if len(results) > 0 { + break + } + } + + return results +} + +func extractToolCallsFromJSONArray(text string) []map[string]interface{} { + trimmed := strings.TrimSpace(text) + if trimmed == "" { + return nil + } + var items []map[string]interface{} + if err := json.Unmarshal([]byte(trimmed), &items); err != nil { + return nil + } + var results []map[string]interface{} + for _, item := range items { + if item == nil { + continue + } + if _, ok := item["name"].(string); ok { + results = append(results, processToolData(item)) + continue + } + if tool, ok := item["tool"].(string); ok { + args := parseToolArgs(item) + if len(args) == 0 { + if params, ok := item["parameters"].(map[string]interface{}); ok { + args = params + } + } + results = append(results, map[string]interface{}{ + "name": tool, + "arguments": args, + }) + } + } + return results +} + +func extractNativeToolCalls(text string) [][2]string { + var results [][2]string + re := regexp.MustCompile(`\[TOOL_CALL:\s*([^\s|]+)\s*\|`) + matches := re.FindAllStringSubmatchIndex(text, -1) + for _, m := range matches { + if len(m) < 4 { + continue + } + name := text[m[2]:m[3]] + jsonStart := m[1] + for jsonStart < len(text) && (text[jsonStart] == ' ' || text[jsonStart] == '\t') { + jsonStart++ + } + var jsonEnd int + inString := false + escapeNext := false + braceDepth := 0 + for i := jsonStart; i < len(text); i++ { + c := text[i] + if escapeNext { + escapeNext = false + continue + } + switch c { + case '\\': + escapeNext = true + case '"': + inString = !inString + case '{', '[': + if !inString { + braceDepth++ + } + case '}', ']': + if !inString { + braceDepth-- + if braceDepth == 0 { + jsonEnd = i + 1 + break + } + } + } + } + if jsonEnd > jsonStart { + results = append(results, [2]string{name, text[jsonStart:jsonEnd]}) + } + } + return results +} + +func normalizeToolCallText(text string) string { + lines := strings.Split(text, "\n") + if len(lines) == 0 { + return text + } + + nonEmpty := 0 + withMargin := 0 + marginRe := regexp.MustCompile(`^\s*│`) + stripRe := regexp.MustCompile(`^\s*│\s?`) + + for _, line := range lines { + if strings.TrimSpace(line) == "" { + continue + } + nonEmpty++ + if marginRe.MatchString(line) { + withMargin++ + } + } + if nonEmpty == 0 || withMargin*2 < nonEmpty { + return text + } + + for i, line := range lines { + lines[i] = stripRe.ReplaceAllString(line, "") + } + return strings.Join(lines, "\n") +} + +func sanitizeMultilineJson(text string) string { + var result strings.Builder + inString := false + escapeNext := false + lines := strings.Split(text, "\n") + + for lineIdx, line := range lines { + for i := 0; i < len(line); i++ { + char := rune(line[i]) + + if escapeNext { + if char == 'n' { + result.WriteString("\\n") + } else { + result.WriteRune('\\') + result.WriteRune(char) + } + escapeNext = false + continue + } + + if char == '\\' && inString { + escapeNext = true + continue + } + + if char == '"' { + if !inString { + inString = true + } else { + backslashCount := 0 + j := i - 1 + for j >= 0 && line[j] == '\\' { + backslashCount++ + j-- + } + if backslashCount%2 == 0 { + inString = false + } + } + } + + result.WriteRune(char) + } + + if inString && lineIdx < len(lines)-1 { + result.WriteString("\\n") + } else if !inString { + result.WriteRune('\n') + } + } + + return result.String() +} + +func extractCodeBlockContent(text string) []string { + var results []string + lines := strings.Split(text, "\n") + inBlock := false + blockStart := -1 + + for i, line := range lines { + if !inBlock { + if strings.Contains(line, "```") { + inBlock = true + blockStart = i + 1 + } + } else { + if strings.TrimSpace(line) == "```" { + content := strings.Join(lines[blockStart:i], "\n") + results = append(results, content) + inBlock = false + blockStart = -1 + } + } + } + return results +} + +func extractJsonObjects(text string) []map[string]interface{} { + var results []map[string]interface{} + startIdx := -1 + braceCount := 0 + + for i, char := range text { + if char == '{' { + if braceCount == 0 { + startIdx = i + } + braceCount++ + } else if char == '}' { + if braceCount > 0 { + braceCount-- + if braceCount == 0 && startIdx != -1 { + potentialJSON := text[startIdx : i+1] + var data map[string]interface{} + if err := json.Unmarshal([]byte(potentialJSON), &data); err == nil { + if _, ok := data["name"].(string); ok { + results = append(results, processToolData(data)) + } + } + } + } + } + } + return results +} + +func processToolData(data map[string]interface{}) map[string]interface{} { + name, _ := data["name"].(string) + args := parseToolArgs(data) + + // Clean up string arguments + for k, v := range args { + if s, ok := v.(string); ok { + s = strings.ReplaceAll(s, "\\u0026", "&") + s = strings.ReplaceAll(s, "&", "&") + s = strings.ReplaceAll(s, "&", "&") + args[k] = s + } + } + return map[string]interface{}{"name": name, "arguments": args} +} + +func parseToolArgs(data map[string]interface{}) map[string]interface{} { + if a, ok := data["arguments"].(map[string]interface{}); ok { + return a + } + if a, ok := data["args"].(map[string]interface{}); ok { + return a + } + if a, ok := data["parameters"].(map[string]interface{}); ok { + return a + } + if a, ok := data["params"].(map[string]interface{}); ok { + return a + } + if raw, ok := data["parameters"].(string); ok && strings.TrimSpace(raw) != "" { + var parsed map[string]interface{} + if err := json.Unmarshal([]byte(raw), &parsed); err == nil { + return parsed + } + } + if raw, ok := data["arguments"].(string); ok && strings.TrimSpace(raw) != "" { + var parsed map[string]interface{} + if err := json.Unmarshal([]byte(raw), &parsed); err == nil { + return parsed + } + } + return make(map[string]interface{}) +} diff --git a/internal/toolcalls/toolcalls_test.go b/internal/toolcalls/toolcalls_test.go new file mode 100644 index 0000000..1b79981 --- /dev/null +++ b/internal/toolcalls/toolcalls_test.go @@ -0,0 +1,48 @@ +package toolcalls + +import "testing" + +func TestExtractAllToolCalls_PlainJSON(t *testing.T) { + in := `I'll do it. + +{ + "name": "run_command", + "arguments": { + "command": "echo hi" + } +}` + + calls := ExtractAllToolCalls(in) + if len(calls) != 1 { + t.Fatalf("expected 1 tool call, got %d", len(calls)) + } + if calls[0]["name"] != "run_command" { + t.Fatalf("expected name=run_command, got %#v", calls[0]["name"]) + } + args, ok := calls[0]["arguments"].(map[string]interface{}) + if !ok { + t.Fatalf("expected arguments map, got %#v", calls[0]["arguments"]) + } + if args["command"] != "echo hi" { + t.Fatalf("expected command=echo hi, got %#v", args["command"]) + } +} + +func TestExtractAllToolCalls_FencedJSON(t *testing.T) { + in := "```json\n{\"name\":\"run_command\",\"arguments\":{\"command\":\"echo hi\"}}\n```" + calls := ExtractAllToolCalls(in) + if len(calls) != 1 { + t.Fatalf("expected 1 tool call, got %d", len(calls)) + } + if calls[0]["name"] != "run_command" { + t.Fatalf("expected name=run_command, got %#v", calls[0]["name"]) + } +} + +func TestExtractAllToolCalls_JSONArray(t *testing.T) { + in := `[{"name":"run_command","arguments":{"command":"echo a"}},{"name":"run_command","arguments":{"command":"echo b"}}]` + calls := ExtractAllToolCalls(in) + if len(calls) != 2 { + t.Fatalf("expected 2 tool calls, got %d", len(calls)) + } +}