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:
ZanzyTHEbar 2026-02-19 12:13:55 +00:00
parent 90c70e8cb0
commit 17634aa9ac
5 changed files with 623 additions and 2 deletions

View file

@ -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()

View file

@ -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

View 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
View 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
View 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
}