From 17634aa9ac4b39fce6b066063f1cddf5efd91ef5 Mon Sep 17 00:00:00 2001 From: ZanzyTHEbar Date: Thu, 19 Feb 2026 12:13:55 +0000 Subject: [PATCH] feat(agent): wire SecureBus into agent loop and main entrypoints SecureBusToolRuntime (pkg/agent/securebus_runtime.go): - Implements fantasy.ToolRuntime wrapping a base runtime - Intercepts Execute calls, dispatches through SecureBus for policy enforcement, secret injection, leak scanning, and audit logging - If SecureBus redacts output (leak detected), replaces base runtime response with the redacted version transparently AgentLoop integration (pkg/agent/loop.go): - secureBus *securebus.Bus field (nil = disabled, zero-overhead path) - SetSecureBus(b) for external injection - SetupSecureBus(ss, cfg) factory: wires capLookup and executor closures against the loop's ToolRegistry, creates and attaches Bus in one call Main entrypoints (cmd/picoclaw/main.go): - setupSecureBus() helper: locates secrets.json in $HOME, selects NoopKeyring or EnvKeyring (when PICOCLAW_MASTER_KEY is set), creates SecretStore, calls agentLoop.SetupSecureBus - Wired into both agentCmd() and gatewayCmd() via defer closeBus() ITR DAG additions (pkg/itr/dag): - planner.go: task graph planning utilities - replan.go: dynamic replanning on tool execution failure --- cmd/picoclaw/main.go | 47 +++++++ pkg/agent/loop.go | 97 +++++++++++++- pkg/agent/securebus_runtime.go | 115 +++++++++++++++++ pkg/itr/dag/planner.go | 227 +++++++++++++++++++++++++++++++++ pkg/itr/dag/replan.go | 139 ++++++++++++++++++++ 5 files changed, 623 insertions(+), 2 deletions(-) create mode 100644 pkg/agent/securebus_runtime.go create mode 100644 pkg/itr/dag/planner.go create mode 100644 pkg/itr/dag/replan.go diff --git a/cmd/picoclaw/main.go b/cmd/picoclaw/main.go index 198ff92c1..45e74e85f 100644 --- a/cmd/picoclaw/main.go +++ b/cmd/picoclaw/main.go @@ -36,6 +36,8 @@ import ( picomemory "github.com/sipeed/picoclaw/pkg/memory" "github.com/sipeed/picoclaw/pkg/memory/delegate" "github.com/sipeed/picoclaw/pkg/migrate" + "github.com/sipeed/picoclaw/pkg/security" + "github.com/sipeed/picoclaw/pkg/security/securebus" "github.com/sipeed/picoclaw/pkg/skills" "github.com/sipeed/picoclaw/pkg/state" "github.com/sipeed/picoclaw/pkg/tools" @@ -420,6 +422,10 @@ func agentCmd() { msgBus := bus.NewMessageBus() agentLoop := agent.NewAgentLoop(cfg, msgBus, languageModel) + // Wire ITR SecureBus ā€” routes all tool calls through capability enforcement. + closeBus := setupSecureBus(agentLoop) + defer closeBus() + // Print agent startup info (only for interactive mode) startupInfo := agentLoop.GetStartupInfo() logger.InfoCF("agent", "Agent initialized", @@ -560,6 +566,10 @@ func gatewayCmd() { msgBus := bus.NewMessageBus() agentLoop := agent.NewAgentLoop(cfg, msgBus, languageModel) + // Wire ITR SecureBus ā€” routes all tool calls through capability enforcement. + closeBus := setupSecureBus(agentLoop) + defer closeBus() + // Print agent startup info fmt.Println("\nšŸ“¦ Agent Status:") startupInfo := agentLoop.GetStartupInfo() @@ -1167,6 +1177,43 @@ func loadConfig() (*config.Config, error) { return config.LoadConfig(getConfigPath()) } +// setupSecureBus wires the Isolated Tool Runtime into an AgentLoop. +// It loads (or lazily creates) the SecretStore from the picoclaw home directory +// and calls agentLoop.SetupSecureBus so all tool calls are routed through +// capability enforcement, secret injection, leak scanning, and audit logging. +// +// If ITR initialisation fails for any non-fatal reason it logs a warning and +// returns without attaching the bus ā€” the agent continues in direct-execution mode. +// The returned closer must be called on shutdown when the bus is non-nil. +func setupSecureBus(agentLoop *agent.AgentLoop) (closer func()) { + home, err := os.UserHomeDir() + if err != nil { + logger.WarnC("itr", "SecureBus: cannot determine home dir ā€” running without ITR") + return func() {} + } + + secretsPath := filepath.Join(home, ".picoclaw", "secrets.json") + + // Use NoopKeyring by default; EnvKeyring when PICOCLAW_MASTER_KEY is set. + var keyring security.KeyringProvider + if mk := os.Getenv("PICOCLAW_MASTER_KEY"); mk != "" { + keyring = security.NewEnvKeyring("PICOCLAW_MASTER_KEY") + } else { + keyring = security.NewNoopKeyring(nil) + } + + ss, err := security.NewSecretStore(secretsPath, keyring) + if err != nil { + logger.WarnCF("itr", "SecureBus: failed to load secret store ā€” running without secret injection", + map[string]interface{}{"error": err.Error()}) + ss = nil + } + + bus := agentLoop.SetupSecureBus(ss, securebus.DefaultBusConfig()) + logger.InfoC("itr", "SecureBus enabled ā€” tool calls routed through ITR") + return bus.Close +} + func cronCmd() { if len(os.Args) < 3 { cronHelp() diff --git a/pkg/agent/loop.go b/pkg/agent/loop.go index b75702144..2a2633cd5 100644 --- a/pkg/agent/loop.go +++ b/pkg/agent/loop.go @@ -31,6 +31,8 @@ import ( "github.com/sipeed/picoclaw/pkg/memory/observation" memstore "github.com/sipeed/picoclaw/pkg/memory/store" "github.com/sipeed/picoclaw/pkg/messages" + "github.com/sipeed/picoclaw/pkg/security" + "github.com/sipeed/picoclaw/pkg/security/securebus" "github.com/sipeed/picoclaw/pkg/session" "github.com/sipeed/picoclaw/pkg/state" "github.com/sipeed/picoclaw/pkg/tools" @@ -51,6 +53,7 @@ type AgentLoop struct { memoryStore *memstore.MemoryStore // 3-tier MemGPT memory (nil if init failed) memDelegate memory.MemoryDelegate // DB delegate (nil if memory disabled) obsManager *observation.Manager // Observational memory (nil if memory disabled) + secureBus *securebus.Bus // ITR SecureBus (nil = disabled, direct execution) activeSessionKey atomic.Value // Current session key for tool access running atomic.Bool summarizing sync.Map // Tracks which sessions are currently being summarized @@ -371,6 +374,35 @@ func (al *AgentLoop) SetChannelManager(cm *channels.Manager) { al.channelManager = cm } +// SetSecureBus attaches a SecureBus to the agent loop. When set, all tool +// calls are routed through the bus for capability enforcement, secret injection, +// leak scanning, and audit logging. Call before the first message is processed. +// Pass nil to disable SecureBus enforcement (direct execution, default). +func (al *AgentLoop) SetSecureBus(b *securebus.Bus) { + al.secureBus = b +} + +// SetupSecureBus creates a SecureBus wired to this loop's tool registry and +// attaches it so all subsequent tool calls are routed through it. +// ss may be nil ā€” secret injection is then disabled but all other enforcement +// (policy, leak scanning, audit) remains active. +// The returned Bus must be closed on shutdown. +func (al *AgentLoop) SetupSecureBus(ss *security.SecretStore, cfg securebus.BusConfig) *securebus.Bus { + capLookup := func(name string) (tools.ToolCapabilities, bool) { + t, ok := al.tools.Get(name) + if !ok { + return tools.ZeroCapabilities(), false + } + return tools.ExtractCapabilities(t), true + } + executor := func(ctx context.Context, name string, args map[string]interface{}) *tools.ToolResult { + return al.tools.Execute(ctx, name, args) + } + b := securebus.New(cfg, ss, capLookup, executor) + al.secureBus = b + return b +} + // RecordLastChannel records the last active channel for this workspace. // This uses the atomic state save mechanism to prevent data loss on crash. func (al *AgentLoop) RecordLastChannel(channel string) error { @@ -628,6 +660,14 @@ func (al *AgentLoop) runAgentLoop(ctx context.Context, opts processOptions) (str if systemPrompt != "" { agentOpts = append(agentOpts, fantasy.WithSystemPrompt(systemPrompt)) } + // Attach SecureBusToolRuntime when ITR is enabled (non-nil bus). + if al.secureBus != nil { + sbrt := SecureBusToolRuntime{ + Bus: al.secureBus, + SessionKey: opts.SessionKey, + } + agentOpts = append(agentOpts, fantasy.WithToolRuntime(sbrt)) + } agent := fantasy.NewAgent(al.languageModel, agentOpts...) logger.DebugCF("agent", "Fantasy agent created", @@ -781,6 +821,14 @@ func (al *AgentLoop) runAgentLoopStreaming(ctx context.Context, opts processOpti if systemPrompt != "" { agentOpts = append(agentOpts, fantasy.WithSystemPrompt(systemPrompt)) } + // Attach SecureBusToolRuntime when ITR is enabled (non-nil bus). + if al.secureBus != nil { + sbrt := SecureBusToolRuntime{ + Bus: al.secureBus, + SessionKey: opts.SessionKey, + } + agentOpts = append(agentOpts, fantasy.WithToolRuntime(sbrt)) + } fantasyAgent := fantasy.NewAgent(al.languageModel, agentOpts...) logger.DebugCF("agent", "Fantasy streaming agent created", @@ -1262,6 +1310,52 @@ func (al *AgentLoop) applyDAGCompression(history []messages.Message) []messages. } // estimateTokens estimates the number of tokens in a message list. +// listConfiguredModels returns a human-readable summary of which providers +// have API credentials configured, and the current default model. +func listConfiguredModels(cfg *config.Config) string { + if cfg == nil { + return "No configuration available." + } + + type entry struct { + name string + key string + } + + // Ordered list of well-known providers. + candidates := []entry{ + {"anthropic", cfg.Providers.Anthropic.APIKey}, + {"openai", cfg.Providers.OpenAI.APIKey}, + {"openrouter", cfg.Providers.OpenRouter.APIKey}, + {"gemini", cfg.Providers.Gemini.APIKey}, + {"groq", cfg.Providers.Groq.APIKey}, + {"zhipu", cfg.Providers.Zhipu.APIKey}, + {"deepseek", cfg.Providers.DeepSeek.APIKey}, + {"moonshot", cfg.Providers.Moonshot.APIKey}, + {"nvidia", cfg.Providers.Nvidia.APIKey}, + {"shengsuanyun", cfg.Providers.ShengSuanYun.APIKey}, + {"vllm", cfg.Providers.VLLM.APIBase}, // vllm uses base URL, not API key + } + + var configured []string + for _, c := range candidates { + if c.key != "" { + configured = append(configured, c.name) + } + } + + current := fmt.Sprintf("Current model: %s", cfg.Agents.Defaults.Model) + if cfg.Agents.Defaults.Provider != "" { + current += fmt.Sprintf(" (provider: %s)", cfg.Agents.Defaults.Provider) + } + + if len(configured) == 0 { + return current + "\nNo providers configured ā€” set API keys in config.json or environment variables." + } + + return current + "\nConfigured providers: " + strings.Join(configured, ", ") +} + func (al *AgentLoop) estimateTokens(msgs []messages.Message) int { totalChars := 0 for _, m := range msgs { @@ -1304,8 +1398,7 @@ func (al *AgentLoop) handleCommand(_ context.Context, msg bus.InboundMessage) (s } switch args[0] { case "models": - // TODO: Fetch available models dynamically if possible - return "Available models: glm-4.7, claude-3-5-sonnet, gpt-4o (configured in config.json/env)", true + return listConfiguredModels(al.cfg), true case "channels": if al.channelManager == nil { return "Channel manager not initialized", true diff --git a/pkg/agent/securebus_runtime.go b/pkg/agent/securebus_runtime.go new file mode 100644 index 000000000..444289a9f --- /dev/null +++ b/pkg/agent/securebus_runtime.go @@ -0,0 +1,115 @@ +package agent + +import ( + "context" + "errors" + + fantasy "charm.land/fantasy" + "github.com/sipeed/picoclaw/pkg/ids" + "github.com/sipeed/picoclaw/pkg/itr" + "github.com/sipeed/picoclaw/pkg/security/securebus" +) + +// SecureBusToolRuntime is a fantasy.ToolRuntime that routes every tool call +// through the SecureBus before (and after) passing it to the underlying runtime. +// +// Pipeline per tool call: +// 1. Serialize tool call args ā†’ ToolRequest +// 2. bus.Execute() ā†’ capability check, secret injection, output scan, audit +// 3. If bus returns a policy error, short-circuit with that error result +// 4. Otherwise delegate to Base runtime for actual execution +// 5. If bus detected a leak, replace Base output with the redacted version +type SecureBusToolRuntime struct { + // Base is the underlying runtime. If nil, the bus result is used directly. + Base fantasy.ToolRuntime + + // Bus is required. + Bus *securebus.Bus + + // SessionKey is forwarded to bus requests for audit tracing. + SessionKey string +} + +// Execute implements fantasy.ToolRuntime. +func (r SecureBusToolRuntime) Execute( + ctx context.Context, + tools []fantasy.AgentTool, + toolCalls []fantasy.ToolCallContent, + onResult func(fantasy.ToolResultContent) error, +) ([]fantasy.ToolResultContent, error) { + if r.Bus == nil || len(toolCalls) == 0 { + if r.Base != nil { + return r.Base.Execute(ctx, tools, toolCalls, onResult) + } + return nil, nil + } + + results := make([]fantasy.ToolResultContent, 0, len(toolCalls)) + + for _, tc := range toolCalls { + reqID := ids.New().String() + req := itr.NewToolExecRequest(reqID, r.SessionKey, tc.ToolCallID, tc.ToolName, tc.Input) + busResp := r.Bus.Execute(ctx, req) + + if busResp.IsError { + // Policy violation or secret resolution failure. + tr := fantasy.ToolResultContent{ + ToolCallID: tc.ToolCallID, + ToolName: tc.ToolName, + Result: fantasy.ToolResultOutputContentError{Error: errors.New(busResp.Result)}, + } + results = append(results, tr) + if onResult != nil { + if err := onResult(tr); err != nil { + return results, err + } + } + continue + } + + // Bus accepted ā€” delegate to Base for actual execution. + if r.Base == nil { + tr := fantasy.ToolResultContent{ + ToolCallID: tc.ToolCallID, + ToolName: tc.ToolName, + Result: fantasy.ToolResultOutputContentText{Text: busResp.Result}, + } + results = append(results, tr) + if onResult != nil { + if err := onResult(tr); err != nil { + return results, err + } + } + continue + } + + // Execute via Base runtime for the single tool call. + baseResults, err := r.Base.Execute(ctx, tools, []fantasy.ToolCallContent{tc}, nil) + if err != nil { + return results, err + } + + for _, br := range baseResults { + if busResp.LeakDetected { + br = overrideResultText(br, busResp.Result) + } + results = append(results, br) + if onResult != nil { + if err := onResult(br); err != nil { + return results, err + } + } + } + } + + return results, nil +} + +// overrideResultText replaces the text output of a ToolResultContent with +// the redacted version produced by the SecureBus. +func overrideResultText(tr fantasy.ToolResultContent, text string) fantasy.ToolResultContent { + if _, ok := tr.Result.(fantasy.ToolResultOutputContentText); ok { + tr.Result = fantasy.ToolResultOutputContentText{Text: text} + } + return tr +} diff --git a/pkg/itr/dag/planner.go b/pkg/itr/dag/planner.go new file mode 100644 index 000000000..eebee59c3 --- /dev/null +++ b/pkg/itr/dag/planner.go @@ -0,0 +1,227 @@ +package dag + +import ( + "context" + "encoding/json" + "fmt" + + "github.com/sipeed/picoclaw/pkg/itr" + "github.com/sipeed/picoclaw/pkg/tools" +) + +// PlannerFunc calls the LLM with a system prompt and user query, returning +// the raw text response and token cost. +type PlannerFunc func(ctx context.Context, systemPrompt, userQuery string) (string, uint32, error) + +// Planner generates DAGPlans from natural-language queries by prompting an +// LLM to produce a structured JSON plan. The LLM outputs a list of nodes +// with dependency edges in a single inference pass (LLMCompiler pattern). +type Planner struct { + callModel PlannerFunc + registry *tools.ToolRegistry + maxParallel uint8 + tokenBudget uint32 +} + +// PlannerConfig configures the DAG planner. +type PlannerConfig struct { + MaxParallel uint8 + TokenBudget uint32 +} + +// DefaultPlannerConfig returns sensible defaults. +func DefaultPlannerConfig() PlannerConfig { + return PlannerConfig{ + MaxParallel: 8, + TokenBudget: 0, + } +} + +// NewPlanner creates a DAG planner. +func NewPlanner(callModel PlannerFunc, registry *tools.ToolRegistry, cfg PlannerConfig) *Planner { + return &Planner{ + callModel: callModel, + registry: registry, + maxParallel: cfg.MaxParallel, + tokenBudget: cfg.TokenBudget, + } +} + +// Plan generates a DAGPlan for the given query. It calls the LLM once to +// produce a structured plan with nodes and dependencies, then validates it. +func (p *Planner) Plan(ctx context.Context, query string, availableTools []string) (*itr.DAGPlan, uint32, error) { + systemPrompt := p.buildSystemPrompt(availableTools) + userPrompt := fmt.Sprintf("Create an execution plan for this task:\n\n%s", query) + + response, tokens, err := p.callModel(ctx, systemPrompt, userPrompt) + if err != nil { + return nil, tokens, fmt.Errorf("planner LLM call failed: %w", err) + } + + plan, err := parsePlanResponse(response) + if err != nil { + return nil, tokens, fmt.Errorf("failed to parse plan: %w", err) + } + + plan.MaxParallel = p.maxParallel + plan.TokenBudget = p.tokenBudget + + if err := validatePlan(plan); err != nil { + return nil, tokens, fmt.Errorf("invalid plan: %w", err) + } + + return plan, tokens, nil +} + +func (p *Planner) buildSystemPrompt(availableTools []string) string { + toolList := "" + if p.registry != nil { + for _, name := range availableTools { + tool, ok := p.registry.Get(name) + if !ok { + continue + } + toolList += fmt.Sprintf("- %s: %s\n", tool.Name(), tool.Description()) + } + } + + return fmt.Sprintf(`You are a task planner that decomposes complex queries into a DAG of tool calls. + +Available tools: +%s +Output ONLY a JSON object with this schema: +{ + "nodes": [ + { + "id": "unique_string", + "type": "tool_exec", + "payload": {"tool_name": "...", "args_json": "{...}"}, + "depends_on": ["other_node_id"] + } + ], + "joiner_query": "Synthesize the results into a final answer for: " +} + +Rules: +- Node IDs must be unique strings. +- Use "#nodeID" in args_json to reference output from another node. +- Nodes with no dependencies run in parallel. +- Keep the plan minimal: prefer fewer nodes with clear dependencies. +- Output valid JSON only, no markdown fences or explanation.`, toolList) +} + +// parsePlanResponse extracts a DAGPlan from the LLM's JSON response. +func parsePlanResponse(response string) (*itr.DAGPlan, error) { + response = extractJSON(response) + + var plan itr.DAGPlan + if err := json.Unmarshal([]byte(response), &plan); err != nil { + return nil, fmt.Errorf("JSON parse error: %w\nraw: %s", err, truncate(response, 500)) + } + + // Unmarshal re-encodes Payload as map[string]interface{} via JSON round-trip. + // Convert payload maps to their proper types. + for i := range plan.Nodes { + n := &plan.Nodes[i] + m, ok := n.Payload.(map[string]interface{}) + if !ok { + continue + } + b, _ := json.Marshal(m) + + switch n.Type { + case itr.CmdToolExec: + var te itr.ToolExec + _ = json.Unmarshal(b, &te) + n.Payload = te + case itr.CmdToolSearch: + var ts itr.ToolSearch + _ = json.Unmarshal(b, &ts) + n.Payload = ts + } + } + + return &plan, nil +} + +// extractJSON strips markdown code fences if present. +func extractJSON(s string) string { + // Strip ```json ... ``` fences + if idx := findIndex(s, "```json"); idx >= 0 { + s = s[idx+7:] + if end := findIndex(s, "```"); end >= 0 { + s = s[:end] + } + } else if idx := findIndex(s, "```"); idx >= 0 { + s = s[idx+3:] + if end := findIndex(s, "```"); end >= 0 { + s = s[:end] + } + } + + // Find first { and last } + start := -1 + end := -1 + for i, c := range s { + if c == '{' { + start = i + break + } + } + for i := len(s) - 1; i >= 0; i-- { + if s[i] == '}' { + end = i + 1 + break + } + } + if start >= 0 && end > start { + return s[start:end] + } + return s +} + +func findIndex(s, substr string) int { + for i := 0; i <= len(s)-len(substr); i++ { + if s[i:i+len(substr)] == substr { + return i + } + } + return -1 +} + +// validatePlan checks structural integrity of a DAGPlan. +func validatePlan(plan *itr.DAGPlan) error { + if len(plan.Nodes) == 0 { + return fmt.Errorf("plan has no nodes") + } + + ids := make(map[string]bool, len(plan.Nodes)) + for _, n := range plan.Nodes { + if n.ID == "" { + return fmt.Errorf("node has empty ID") + } + if ids[n.ID] { + return fmt.Errorf("duplicate node ID: %s", n.ID) + } + ids[n.ID] = true + } + + for _, n := range plan.Nodes { + for _, dep := range n.DependsOn { + if !ids[dep] { + return fmt.Errorf("node %s depends on unknown node %s", n.ID, dep) + } + if dep == n.ID { + return fmt.Errorf("node %s depends on itself", n.ID) + } + } + } + + // Verify no cycles using topological sort on temporary nodeStates. + states := make(map[string]*nodeState, len(plan.Nodes)) + for _, n := range plan.Nodes { + states[n.ID] = newNodeState(n.ID, n.DependsOn) + } + _, err := topologicalOrder(states) + return err +} diff --git a/pkg/itr/dag/replan.go b/pkg/itr/dag/replan.go new file mode 100644 index 000000000..5f4dbcf41 --- /dev/null +++ b/pkg/itr/dag/replan.go @@ -0,0 +1,139 @@ +package dag + +import ( + "context" + "fmt" + + "github.com/sipeed/picoclaw/pkg/itr" +) + +// ReplanConfig configures the iterative replanning loop. +type ReplanConfig struct { + MaxReplans int // maximum replanning iterations; 0 ā†’ 3 +} + +// DefaultReplanConfig returns sensible defaults. +func DefaultReplanConfig() ReplanConfig { + return ReplanConfig{MaxReplans: 3} +} + +// ReplanResult holds the combined output of an iterative DAG execution with +// optional replanning passes. +type ReplanResult struct { + FinalAnswer string + TotalTokens uint32 + Iterations int +} + +// ReplanLoop runs the DAG executor with iterative replanning. After each +// execution, the Joiner evaluates whether the answer is complete. If not, +// it triggers another planning pass with the previous results as context. +// +// The loop terminates when: +// - The Joiner's answer does not contain the replan sentinel, or +// - MaxReplans iterations are exhausted, or +// - An error occurs. +func ReplanLoop( + ctx context.Context, + executor *Executor, + planner *Planner, + sessionKey string, + query string, + availableTools []string, + cfg ReplanConfig, +) (*ReplanResult, error) { + maxReplans := cfg.MaxReplans + if maxReplans <= 0 { + maxReplans = 3 + } + + var totalTokens uint32 + previousResults := "" + + for i := 0; i <= maxReplans; i++ { + planQuery := query + if previousResults != "" { + planQuery = fmt.Sprintf( + "Previous execution produced partial results:\n%s\n\nOriginal task: %s\n\nPlan additional steps to complete the task.", + truncate(previousResults, 4000), query) + } + + plan, planTokens, err := planner.Plan(ctx, planQuery, availableTools) + if err != nil { + return nil, fmt.Errorf("replan iteration %d: planning failed: %w", i, err) + } + totalTokens += planTokens + + result, err := executor.Execute(ctx, sessionKey, plan) + if err != nil { + return nil, fmt.Errorf("replan iteration %d: execution failed: %w", i, err) + } + totalTokens += result.TotalTokens + + if !needsReplan(result.FinalAnswer) { + return &ReplanResult{ + FinalAnswer: result.FinalAnswer, + TotalTokens: totalTokens, + Iterations: i + 1, + }, nil + } + + previousResults = result.FinalAnswer + } + + return &ReplanResult{ + FinalAnswer: previousResults, + TotalTokens: totalTokens, + Iterations: maxReplans + 1, + }, nil +} + +const replanSentinel = "[NEEDS_MORE_STEPS]" + +// needsReplan checks whether the Joiner's output indicates the task is +// incomplete and requires another planning pass. +func needsReplan(answer string) bool { + return findIndex(answer, replanSentinel) >= 0 +} + +// ExecuteWithReplan is a convenience wrapper that runs a pre-built DAGPlan +// through the executor, then optionally replans if the Joiner signals +// incompleteness. +func ExecuteWithReplan( + ctx context.Context, + executor *Executor, + planner *Planner, + sessionKey string, + initialPlan *itr.DAGPlan, + query string, + availableTools []string, + cfg ReplanConfig, +) (*ReplanResult, error) { + maxReplans := cfg.MaxReplans + if maxReplans <= 0 { + maxReplans = 3 + } + + var totalTokens uint32 + + result, err := executor.Execute(ctx, sessionKey, initialPlan) + if err != nil { + return nil, err + } + totalTokens += result.TotalTokens + + if !needsReplan(result.FinalAnswer) { + return &ReplanResult{ + FinalAnswer: result.FinalAnswer, + TotalTokens: totalTokens, + Iterations: 1, + }, nil + } + + replanResult, err := ReplanLoop(ctx, executor, planner, sessionKey, query, availableTools, cfg) + if err != nil { + return nil, err + } + replanResult.TotalTokens += totalTokens + return replanResult, nil +}