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
This commit is contained in:
parent
90c70e8cb0
commit
17634aa9ac
5 changed files with 623 additions and 2 deletions
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
115
pkg/agent/securebus_runtime.go
Normal file
115
pkg/agent/securebus_runtime.go
Normal file
|
|
@ -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
|
||||
}
|
||||
227
pkg/itr/dag/planner.go
Normal file
227
pkg/itr/dag/planner.go
Normal file
|
|
@ -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: <original query>"
|
||||
}
|
||||
|
||||
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
|
||||
}
|
||||
139
pkg/itr/dag/replan.go
Normal file
139
pkg/itr/dag/replan.go
Normal file
|
|
@ -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
|
||||
}
|
||||
Loading…
Add table
Reference in a new issue