97 lines
2.5 KiB
Go
97 lines
2.5 KiB
Go
package behavior
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"sync"
|
|
|
|
"github.com/sipeed/picoclaw/pkg/agent"
|
|
)
|
|
|
|
type turnStats struct {
|
|
toolCalls int
|
|
totalBytes int64
|
|
}
|
|
|
|
// Monitor implements agent.ToolInterceptor and agent.EventObserver to detect behavioral anomalies.
|
|
type Monitor struct {
|
|
MaxToolCalls int
|
|
MaxTotalBytes int64
|
|
|
|
mu sync.Mutex
|
|
turns map[string]*turnStats
|
|
}
|
|
|
|
// Ensure Monitor implements necessary interfaces.
|
|
var _ agent.ToolInterceptor = (*Monitor)(nil)
|
|
var _ agent.EventObserver = (*Monitor)(nil)
|
|
|
|
// NewMonitor creates a new behavioral monitor.
|
|
func NewMonitor(maxCalls int, maxBytes int64) *Monitor {
|
|
return &Monitor{
|
|
MaxToolCalls: maxCalls,
|
|
MaxTotalBytes: maxBytes,
|
|
turns: make(map[string]*turnStats),
|
|
}
|
|
}
|
|
|
|
func (m *Monitor) OnEvent(ctx context.Context, evt agent.Event) error {
|
|
if evt.Kind == agent.EventKindTurnEnd {
|
|
m.mu.Lock()
|
|
delete(m.turns, evt.Meta.TurnID)
|
|
m.mu.Unlock()
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (m *Monitor) BeforeTool(ctx context.Context, call *agent.ToolCallHookRequest) (*agent.ToolCallHookRequest, agent.HookDecision, error) {
|
|
if call == nil {
|
|
return nil, agent.HookDecision{}, nil
|
|
}
|
|
|
|
m.mu.Lock()
|
|
defer m.mu.Unlock()
|
|
|
|
stats, ok := m.turns[call.Meta.TurnID]
|
|
if !ok {
|
|
stats = &turnStats{}
|
|
m.turns[call.Meta.TurnID] = stats
|
|
}
|
|
|
|
stats.toolCalls++
|
|
|
|
if m.MaxToolCalls > 0 && stats.toolCalls > m.MaxToolCalls {
|
|
return call, agent.HookDecision{
|
|
Action: agent.HookActionAbortTurn,
|
|
Reason: fmt.Sprintf("Behavioral defense: Tool call limit (%d) exceeded in a single turn", m.MaxToolCalls),
|
|
}, nil
|
|
}
|
|
|
|
return call, agent.HookDecision{Action: agent.HookActionContinue}, nil
|
|
}
|
|
|
|
func (m *Monitor) AfterTool(ctx context.Context, resp *agent.ToolResultHookResponse) (*agent.ToolResultHookResponse, agent.HookDecision, error) {
|
|
if resp == nil || resp.Result == nil {
|
|
return resp, agent.HookDecision{Action: agent.HookActionContinue}, nil
|
|
}
|
|
|
|
m.mu.Lock()
|
|
defer m.mu.Unlock()
|
|
|
|
stats, ok := m.turns[resp.Meta.TurnID]
|
|
if !ok {
|
|
// Should have been created in BeforeTool, but handle just in case.
|
|
return resp, agent.HookDecision{Action: agent.HookActionContinue}, nil
|
|
}
|
|
|
|
stats.totalBytes += int64(len(resp.Result.ForLLM))
|
|
|
|
if m.MaxTotalBytes > 0 && stats.totalBytes > m.MaxTotalBytes {
|
|
return resp, agent.HookDecision{
|
|
Action: agent.HookActionAbortTurn,
|
|
Reason: fmt.Sprintf("Behavioral defense: Cumulative tool output size limit (%d bytes) exceeded in a single turn", m.MaxTotalBytes),
|
|
}, nil
|
|
}
|
|
|
|
return resp, agent.HookDecision{Action: agent.HookActionContinue}, nil
|
|
}
|