temp
This commit is contained in:
parent
b26337501c
commit
fe1be76bc6
13 changed files with 4633 additions and 19 deletions
369
pkg/agent/artifact.go
Normal file
369
pkg/agent/artifact.go
Normal file
|
|
@ -0,0 +1,369 @@
|
||||||
|
package agent
|
||||||
|
|
||||||
|
import (
|
||||||
|
"crypto/sha256"
|
||||||
|
"encoding/hex"
|
||||||
|
"fmt"
|
||||||
|
"sync"
|
||||||
|
"time"
|
||||||
|
)
|
||||||
|
|
||||||
|
// Artifact represents a trackable product generated by the LLM
|
||||||
|
// such as code, images, documents, or data files.
|
||||||
|
type Artifact struct {
|
||||||
|
ID string `json:"id"` // Unique identifier
|
||||||
|
Type ArtifactType `json:"type"` // Type of artifact
|
||||||
|
Content string `json:"content"` // Artifact content
|
||||||
|
MIME string `json:"mime"` // MIME type if applicable
|
||||||
|
Size int64 `json:"size"` // Size in bytes
|
||||||
|
Created time.Time `json:"created"` // Creation timestamp
|
||||||
|
SessionID string `json:"session_id"` // Associated session
|
||||||
|
MessageIndex int `json:"message_index"` // Index in message history
|
||||||
|
Metadata map[string]any `json:"metadata,omitempty"` // Additional metadata
|
||||||
|
|
||||||
|
// Optional fields for enhanced tracking
|
||||||
|
Language string `json:"language,omitempty"` // Programming language (for code)
|
||||||
|
Filename string `json:"filename,omitempty"` // Suggested filename
|
||||||
|
Description string `json:"description,omitempty"` // Human-readable description
|
||||||
|
}
|
||||||
|
|
||||||
|
// ArtifactManager manages the lifecycle of artifacts within agent sessions
|
||||||
|
type ArtifactManager struct {
|
||||||
|
mu sync.RWMutex
|
||||||
|
artifacts map[string]*Artifact // ID -> Artifact
|
||||||
|
bySession map[string][]string // SessionID -> []ArtifactID
|
||||||
|
}
|
||||||
|
|
||||||
|
// NewArtifactManager creates a new artifact manager
|
||||||
|
func NewArtifactManager() *ArtifactManager {
|
||||||
|
return &ArtifactManager{
|
||||||
|
artifacts: make(map[string]*Artifact),
|
||||||
|
bySession: make(map[string][]string),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// CreateArtifact creates and registers a new artifact
|
||||||
|
func (am *ArtifactManager) CreateArtifact(artifactType ArtifactType, content, sessionID string) *Artifact {
|
||||||
|
artifact := &Artifact{
|
||||||
|
ID: generateArtifactID(content),
|
||||||
|
Type: artifactType,
|
||||||
|
Content: content,
|
||||||
|
Size: int64(len(content)),
|
||||||
|
Created: time.Now(),
|
||||||
|
SessionID: sessionID,
|
||||||
|
Metadata: make(map[string]any),
|
||||||
|
}
|
||||||
|
|
||||||
|
am.mu.Lock()
|
||||||
|
defer am.mu.Unlock()
|
||||||
|
|
||||||
|
am.artifacts[artifact.ID] = artifact
|
||||||
|
am.bySession[sessionID] = append(am.bySession[sessionID], artifact.ID)
|
||||||
|
|
||||||
|
return artifact
|
||||||
|
}
|
||||||
|
|
||||||
|
// AddArtifact registers an existing artifact
|
||||||
|
func (am *ArtifactManager) AddArtifact(artifact *Artifact) {
|
||||||
|
am.mu.Lock()
|
||||||
|
defer am.mu.Unlock()
|
||||||
|
|
||||||
|
if artifact.ID == "" {
|
||||||
|
artifact.ID = generateArtifactID(artifact.Content)
|
||||||
|
}
|
||||||
|
|
||||||
|
am.artifacts[artifact.ID] = artifact
|
||||||
|
am.bySession[artifact.SessionID] = append(am.bySession[artifact.SessionID], artifact.ID)
|
||||||
|
}
|
||||||
|
|
||||||
|
// GetArtifact retrieves an artifact by ID
|
||||||
|
func (am *ArtifactManager) GetArtifact(id string) (*Artifact, bool) {
|
||||||
|
am.mu.RLock()
|
||||||
|
defer am.mu.RUnlock()
|
||||||
|
|
||||||
|
artifact, ok := am.artifacts[id]
|
||||||
|
return artifact, ok
|
||||||
|
}
|
||||||
|
|
||||||
|
// GetSessionArtifacts retrieves all artifacts for a session
|
||||||
|
func (am *ArtifactManager) GetSessionArtifacts(sessionID string) []*Artifact {
|
||||||
|
am.mu.RLock()
|
||||||
|
defer am.mu.RUnlock()
|
||||||
|
|
||||||
|
ids, ok := am.bySession[sessionID]
|
||||||
|
if !ok {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
artifacts := make([]*Artifact, 0, len(ids))
|
||||||
|
for _, id := range ids {
|
||||||
|
if artifact, ok := am.artifacts[id]; ok {
|
||||||
|
artifacts = append(artifacts, artifact)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
return artifacts
|
||||||
|
}
|
||||||
|
|
||||||
|
// GetRecentArtifacts retrieves the N most recent artifacts for a session
|
||||||
|
func (am *ArtifactManager) GetRecentArtifacts(sessionID string, limit int) []*Artifact {
|
||||||
|
artifacts := am.GetSessionArtifacts(sessionID)
|
||||||
|
|
||||||
|
if len(artifacts) <= limit {
|
||||||
|
return artifacts
|
||||||
|
}
|
||||||
|
|
||||||
|
// Return the last N artifacts
|
||||||
|
return artifacts[len(artifacts)-limit:]
|
||||||
|
}
|
||||||
|
|
||||||
|
// GetArtifactsByType retrieves all artifacts of a specific type in a session
|
||||||
|
func (am *ArtifactManager) GetArtifactsByType(sessionID string, artifactType ArtifactType) []*Artifact {
|
||||||
|
am.mu.RLock()
|
||||||
|
defer am.mu.RUnlock()
|
||||||
|
|
||||||
|
ids, ok := am.bySession[sessionID]
|
||||||
|
if !ok {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
var result []*Artifact
|
||||||
|
for _, id := range ids {
|
||||||
|
if artifact, ok := am.artifacts[id]; ok && artifact.Type == artifactType {
|
||||||
|
result = append(result, artifact)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
return result
|
||||||
|
}
|
||||||
|
|
||||||
|
// UpdateArtifact updates an existing artifact
|
||||||
|
func (am *ArtifactManager) UpdateArtifact(id string, updateFn func(*Artifact)) error {
|
||||||
|
am.mu.Lock()
|
||||||
|
defer am.mu.Unlock()
|
||||||
|
|
||||||
|
artifact, ok := am.artifacts[id]
|
||||||
|
if !ok {
|
||||||
|
return fmt.Errorf("artifact not found: %s", id)
|
||||||
|
}
|
||||||
|
|
||||||
|
updateFn(artifact)
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// DeleteArtifact removes an artifact
|
||||||
|
func (am *ArtifactManager) DeleteArtifact(id string) error {
|
||||||
|
am.mu.Lock()
|
||||||
|
defer am.mu.Unlock()
|
||||||
|
|
||||||
|
artifact, ok := am.artifacts[id]
|
||||||
|
if !ok {
|
||||||
|
return fmt.Errorf("artifact not found: %s", id)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Remove from artifacts map
|
||||||
|
delete(am.artifacts, id)
|
||||||
|
|
||||||
|
// Remove from session index
|
||||||
|
if ids, ok := am.bySession[artifact.SessionID]; ok {
|
||||||
|
newIds := make([]string, 0, len(ids)-1)
|
||||||
|
for _, artifactID := range ids {
|
||||||
|
if artifactID != id {
|
||||||
|
newIds = append(newIds, artifactID)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
am.bySession[artifact.SessionID] = newIds
|
||||||
|
}
|
||||||
|
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// DeleteSessionArtifacts removes all artifacts for a session
|
||||||
|
func (am *ArtifactManager) DeleteSessionArtifacts(sessionID string) int {
|
||||||
|
am.mu.Lock()
|
||||||
|
defer am.mu.Unlock()
|
||||||
|
|
||||||
|
ids, ok := am.bySession[sessionID]
|
||||||
|
if !ok {
|
||||||
|
return 0
|
||||||
|
}
|
||||||
|
|
||||||
|
count := 0
|
||||||
|
for _, id := range ids {
|
||||||
|
delete(am.artifacts, id)
|
||||||
|
count++
|
||||||
|
}
|
||||||
|
|
||||||
|
delete(am.bySession, sessionID)
|
||||||
|
return count
|
||||||
|
}
|
||||||
|
|
||||||
|
// CountArtifacts returns the total number of artifacts
|
||||||
|
func (am *ArtifactManager) CountArtifacts() int {
|
||||||
|
am.mu.RLock()
|
||||||
|
defer am.mu.RUnlock()
|
||||||
|
return len(am.artifacts)
|
||||||
|
}
|
||||||
|
|
||||||
|
// CountSessionArtifacts returns the number of artifacts in a session
|
||||||
|
func (am *ArtifactManager) CountSessionArtifacts(sessionID string) int {
|
||||||
|
am.mu.RLock()
|
||||||
|
defer am.mu.RUnlock()
|
||||||
|
return len(am.bySession[sessionID])
|
||||||
|
}
|
||||||
|
|
||||||
|
// GetStats returns statistics about artifacts
|
||||||
|
func (am *ArtifactManager) GetStats() ArtifactStats {
|
||||||
|
am.mu.RLock()
|
||||||
|
defer am.mu.RUnlock()
|
||||||
|
|
||||||
|
stats := ArtifactStats{
|
||||||
|
Total: len(am.artifacts),
|
||||||
|
ByType: make(map[ArtifactType]int),
|
||||||
|
BySession: make(map[string]int),
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, artifact := range am.artifacts {
|
||||||
|
stats.ByType[artifact.Type]++
|
||||||
|
stats.TotalSize += artifact.Size
|
||||||
|
}
|
||||||
|
|
||||||
|
for sessionID, ids := range am.bySession {
|
||||||
|
stats.BySession[sessionID] = len(ids)
|
||||||
|
}
|
||||||
|
|
||||||
|
return stats
|
||||||
|
}
|
||||||
|
|
||||||
|
// ArtifactStats provides statistics about artifacts
|
||||||
|
type ArtifactStats struct {
|
||||||
|
Total int
|
||||||
|
TotalSize int64
|
||||||
|
ByType map[ArtifactType]int
|
||||||
|
BySession map[string]int
|
||||||
|
}
|
||||||
|
|
||||||
|
// String returns a human-readable representation of the stats
|
||||||
|
func (s ArtifactStats) String() string {
|
||||||
|
return fmt.Sprintf(
|
||||||
|
"Artifacts: %d total, %s total size, %d sessions",
|
||||||
|
s.Total,
|
||||||
|
formatBytes(s.TotalSize),
|
||||||
|
len(s.BySession),
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
// generateArtifactID generates a unique ID for an artifact based on its content
|
||||||
|
func generateArtifactID(content string) string {
|
||||||
|
hash := sha256.Sum256([]byte(content))
|
||||||
|
return "artifact_" + hex.EncodeToString(hash[:8])
|
||||||
|
}
|
||||||
|
|
||||||
|
// ArtifactFilter defines a function type for filtering artifacts
|
||||||
|
type ArtifactFilter func(*Artifact) bool
|
||||||
|
|
||||||
|
// FilterArtifacts returns artifacts that match the given filter
|
||||||
|
func FilterArtifacts(artifacts []*Artifact, filter ArtifactFilter) []*Artifact {
|
||||||
|
var result []*Artifact
|
||||||
|
for _, artifact := range artifacts {
|
||||||
|
if filter(artifact) {
|
||||||
|
result = append(result, artifact)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return result
|
||||||
|
}
|
||||||
|
|
||||||
|
// Common artifact filters
|
||||||
|
|
||||||
|
// FilterArtifactsByType returns a filter that matches artifacts of the given type
|
||||||
|
func FilterArtifactsByType(artifactType ArtifactType) ArtifactFilter {
|
||||||
|
return func(a *Artifact) bool {
|
||||||
|
return a.Type == artifactType
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// FilterArtifactsByLanguage returns a filter that matches artifacts with the given language
|
||||||
|
func FilterArtifactsByLanguage(language string) ArtifactFilter {
|
||||||
|
return func(a *Artifact) bool {
|
||||||
|
return a.Language == language
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// FilterArtifactsByMinSize returns a filter that matches artifacts above a size threshold
|
||||||
|
func FilterArtifactsByMinSize(minSize int64) ArtifactFilter {
|
||||||
|
return func(a *Artifact) bool {
|
||||||
|
return a.Size >= minSize
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// FilterArtifactsSince returns a filter that matches artifacts created after a timestamp
|
||||||
|
func FilterArtifactsSince(since time.Time) ArtifactFilter {
|
||||||
|
return func(a *Artifact) bool {
|
||||||
|
return a.Created.After(since)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Helper methods for Artifact
|
||||||
|
|
||||||
|
// WithMetadata adds metadata to an artifact (builder pattern)
|
||||||
|
func (a *Artifact) WithMetadata(key string, value any) *Artifact {
|
||||||
|
if a.Metadata == nil {
|
||||||
|
a.Metadata = make(map[string]any)
|
||||||
|
}
|
||||||
|
a.Metadata[key] = value
|
||||||
|
return a
|
||||||
|
}
|
||||||
|
|
||||||
|
// WithLanguage sets the programming language (builder pattern)
|
||||||
|
func (a *Artifact) WithLanguage(language string) *Artifact {
|
||||||
|
a.Language = language
|
||||||
|
return a
|
||||||
|
}
|
||||||
|
|
||||||
|
// WithFilename sets the suggested filename (builder pattern)
|
||||||
|
func (a *Artifact) WithFilename(filename string) *Artifact {
|
||||||
|
a.Filename = filename
|
||||||
|
return a
|
||||||
|
}
|
||||||
|
|
||||||
|
// WithDescription sets the description (builder pattern)
|
||||||
|
func (a *Artifact) WithDescription(description string) *Artifact {
|
||||||
|
a.Description = description
|
||||||
|
return a
|
||||||
|
}
|
||||||
|
|
||||||
|
// WithMIME sets the MIME type (builder pattern)
|
||||||
|
func (a *Artifact) WithMIME(mime string) *Artifact {
|
||||||
|
a.MIME = mime
|
||||||
|
return a
|
||||||
|
}
|
||||||
|
|
||||||
|
// ToMessage converts an artifact to an AgentMessage
|
||||||
|
func (a *Artifact) ToMessage() *AgentMessage {
|
||||||
|
return &AgentMessage{
|
||||||
|
Role: "assistant",
|
||||||
|
Content: a.Content,
|
||||||
|
Type: MessageTypeArtifact,
|
||||||
|
ArtifactID: a.ID,
|
||||||
|
ArtifactType: a.Type,
|
||||||
|
ArtifactMIME: a.MIME,
|
||||||
|
ArtifactSize: a.Size,
|
||||||
|
SessionID: a.SessionID,
|
||||||
|
Timestamp: a.Created,
|
||||||
|
Metadata: a.Metadata,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Clone creates a deep copy of the artifact
|
||||||
|
func (a *Artifact) Clone() *Artifact {
|
||||||
|
clone := *a
|
||||||
|
|
||||||
|
if a.Metadata != nil {
|
||||||
|
clone.Metadata = make(map[string]any, len(a.Metadata))
|
||||||
|
for k, v := range a.Metadata {
|
||||||
|
clone.Metadata[k] = v
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
return &clone
|
||||||
|
}
|
||||||
458
pkg/agent/artifact_test.go
Normal file
458
pkg/agent/artifact_test.go
Normal file
|
|
@ -0,0 +1,458 @@
|
||||||
|
package agent
|
||||||
|
|
||||||
|
import (
|
||||||
|
"testing"
|
||||||
|
"time"
|
||||||
|
)
|
||||||
|
|
||||||
|
// TestArtifact tests basic Artifact functionality
|
||||||
|
func TestArtifact(t *testing.T) {
|
||||||
|
t.Run("ToMessage", func(t *testing.T) {
|
||||||
|
artifact := &Artifact{
|
||||||
|
ID: "artifact_123",
|
||||||
|
Type: ArtifactTypeCode,
|
||||||
|
Content: "package main",
|
||||||
|
Size: 12,
|
||||||
|
Created: time.Now(),
|
||||||
|
SessionID: "session_1",
|
||||||
|
}
|
||||||
|
|
||||||
|
msg := artifact.ToMessage()
|
||||||
|
|
||||||
|
if msg.Type != MessageTypeArtifact {
|
||||||
|
t.Errorf("Expected type 'artifact', got '%s'", msg.Type)
|
||||||
|
}
|
||||||
|
if msg.ArtifactID != artifact.ID {
|
||||||
|
t.Error("Artifact ID mismatch")
|
||||||
|
}
|
||||||
|
if msg.Content != artifact.Content {
|
||||||
|
t.Error("Content mismatch")
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("BuilderPattern", func(t *testing.T) {
|
||||||
|
artifact := &Artifact{ID: "test"}
|
||||||
|
artifact.
|
||||||
|
WithLanguage("go").
|
||||||
|
WithFilename("main.go").
|
||||||
|
WithDescription("Main entry point").
|
||||||
|
WithMIME("text/x-go").
|
||||||
|
WithMetadata("version", "1.0")
|
||||||
|
|
||||||
|
if artifact.Language != "go" {
|
||||||
|
t.Error("Language not set")
|
||||||
|
}
|
||||||
|
if artifact.Filename != "main.go" {
|
||||||
|
t.Error("Filename not set")
|
||||||
|
}
|
||||||
|
if artifact.Description != "Main entry point" {
|
||||||
|
t.Error("Description not set")
|
||||||
|
}
|
||||||
|
if artifact.MIME != "text/x-go" {
|
||||||
|
t.Error("MIME not set")
|
||||||
|
}
|
||||||
|
if artifact.Metadata["version"] != "1.0" {
|
||||||
|
t.Error("Metadata not set")
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("Clone", func(t *testing.T) {
|
||||||
|
original := &Artifact{
|
||||||
|
ID: "test",
|
||||||
|
Content: "original",
|
||||||
|
Metadata: map[string]any{"key": "value"},
|
||||||
|
}
|
||||||
|
|
||||||
|
cloned := original.Clone()
|
||||||
|
cloned.Content = "modified"
|
||||||
|
cloned.Metadata["key"] = "modified"
|
||||||
|
|
||||||
|
if original.Content != "original" {
|
||||||
|
t.Error("Original content was modified")
|
||||||
|
}
|
||||||
|
if original.Metadata["key"] != "value" {
|
||||||
|
t.Error("Original metadata was modified")
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
// TestArtifactManager tests ArtifactManager functionality
|
||||||
|
func TestArtifactManager(t *testing.T) {
|
||||||
|
t.Run("CreateArtifact", func(t *testing.T) {
|
||||||
|
am := NewArtifactManager()
|
||||||
|
artifact := am.CreateArtifact(ArtifactTypeCode, "package main", "session_1")
|
||||||
|
|
||||||
|
if artifact.ID == "" {
|
||||||
|
t.Error("Expected artifact to have an ID")
|
||||||
|
}
|
||||||
|
if artifact.Type != ArtifactTypeCode {
|
||||||
|
t.Error("Artifact type mismatch")
|
||||||
|
}
|
||||||
|
if artifact.SessionID != "session_1" {
|
||||||
|
t.Error("Session ID mismatch")
|
||||||
|
}
|
||||||
|
if artifact.Size != 12 {
|
||||||
|
t.Errorf("Expected size 12, got %d", artifact.Size)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Verify it's stored
|
||||||
|
retrieved, ok := am.GetArtifact(artifact.ID)
|
||||||
|
if !ok {
|
||||||
|
t.Error("Artifact not found in manager")
|
||||||
|
}
|
||||||
|
if retrieved.ID != artifact.ID {
|
||||||
|
t.Error("Retrieved artifact ID mismatch")
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("AddArtifact", func(t *testing.T) {
|
||||||
|
am := NewArtifactManager()
|
||||||
|
artifact := &Artifact{
|
||||||
|
ID: "custom_id",
|
||||||
|
Type: ArtifactTypeImage,
|
||||||
|
Content: "image data",
|
||||||
|
SessionID: "session_1",
|
||||||
|
}
|
||||||
|
|
||||||
|
am.AddArtifact(artifact)
|
||||||
|
|
||||||
|
retrieved, ok := am.GetArtifact("custom_id")
|
||||||
|
if !ok {
|
||||||
|
t.Error("Artifact not found")
|
||||||
|
}
|
||||||
|
if retrieved.Type != ArtifactTypeImage {
|
||||||
|
t.Error("Type mismatch")
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("GetSessionArtifacts", func(t *testing.T) {
|
||||||
|
am := NewArtifactManager()
|
||||||
|
am.CreateArtifact(ArtifactTypeCode, "code1", "session_1")
|
||||||
|
am.CreateArtifact(ArtifactTypeCode, "code2", "session_1")
|
||||||
|
am.CreateArtifact(ArtifactTypeCode, "code3", "session_2")
|
||||||
|
|
||||||
|
artifacts := am.GetSessionArtifacts("session_1")
|
||||||
|
if len(artifacts) != 2 {
|
||||||
|
t.Errorf("Expected 2 artifacts for session_1, got %d", len(artifacts))
|
||||||
|
}
|
||||||
|
|
||||||
|
artifacts = am.GetSessionArtifacts("session_2")
|
||||||
|
if len(artifacts) != 1 {
|
||||||
|
t.Errorf("Expected 1 artifact for session_2, got %d", len(artifacts))
|
||||||
|
}
|
||||||
|
|
||||||
|
artifacts = am.GetSessionArtifacts("nonexistent")
|
||||||
|
if artifacts != nil {
|
||||||
|
t.Error("Expected nil for nonexistent session")
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("GetRecentArtifacts", func(t *testing.T) {
|
||||||
|
am := NewArtifactManager()
|
||||||
|
for i := 0; i < 5; i++ {
|
||||||
|
am.CreateArtifact(ArtifactTypeCode, "code", "session_1")
|
||||||
|
time.Sleep(time.Millisecond) // Ensure different timestamps
|
||||||
|
}
|
||||||
|
|
||||||
|
recent := am.GetRecentArtifacts("session_1", 3)
|
||||||
|
if len(recent) != 3 {
|
||||||
|
t.Errorf("Expected 3 recent artifacts, got %d", len(recent))
|
||||||
|
}
|
||||||
|
|
||||||
|
// Should return all if limit exceeds count
|
||||||
|
recent = am.GetRecentArtifacts("session_1", 10)
|
||||||
|
if len(recent) != 5 {
|
||||||
|
t.Errorf("Expected 5 artifacts, got %d", len(recent))
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("GetArtifactsByType", func(t *testing.T) {
|
||||||
|
am := NewArtifactManager()
|
||||||
|
am.CreateArtifact(ArtifactTypeCode, "code", "session_1")
|
||||||
|
am.CreateArtifact(ArtifactTypeImage, "image", "session_1")
|
||||||
|
am.CreateArtifact(ArtifactTypeCode, "code2", "session_1")
|
||||||
|
|
||||||
|
codeArtifacts := am.GetArtifactsByType("session_1", ArtifactTypeCode)
|
||||||
|
if len(codeArtifacts) != 2 {
|
||||||
|
t.Errorf("Expected 2 code artifacts, got %d", len(codeArtifacts))
|
||||||
|
}
|
||||||
|
|
||||||
|
imageArtifacts := am.GetArtifactsByType("session_1", ArtifactTypeImage)
|
||||||
|
if len(imageArtifacts) != 1 {
|
||||||
|
t.Errorf("Expected 1 image artifact, got %d", len(imageArtifacts))
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("UpdateArtifact", func(t *testing.T) {
|
||||||
|
am := NewArtifactManager()
|
||||||
|
artifact := am.CreateArtifact(ArtifactTypeCode, "original", "session_1")
|
||||||
|
|
||||||
|
err := am.UpdateArtifact(artifact.ID, func(a *Artifact) {
|
||||||
|
a.Content = "updated"
|
||||||
|
a.Language = "go"
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
t.Errorf("Update failed: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
updated, _ := am.GetArtifact(artifact.ID)
|
||||||
|
if updated.Content != "updated" {
|
||||||
|
t.Error("Content not updated")
|
||||||
|
}
|
||||||
|
if updated.Language != "go" {
|
||||||
|
t.Error("Language not updated")
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("UpdateArtifact_NotFound", func(t *testing.T) {
|
||||||
|
am := NewArtifactManager()
|
||||||
|
err := am.UpdateArtifact("nonexistent", func(a *Artifact) {})
|
||||||
|
|
||||||
|
if err == nil {
|
||||||
|
t.Error("Expected error for nonexistent artifact")
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("DeleteArtifact", func(t *testing.T) {
|
||||||
|
am := NewArtifactManager()
|
||||||
|
artifact := am.CreateArtifact(ArtifactTypeCode, "code", "session_1")
|
||||||
|
|
||||||
|
err := am.DeleteArtifact(artifact.ID)
|
||||||
|
if err != nil {
|
||||||
|
t.Errorf("Delete failed: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
_, ok := am.GetArtifact(artifact.ID)
|
||||||
|
if ok {
|
||||||
|
t.Error("Artifact still exists after deletion")
|
||||||
|
}
|
||||||
|
|
||||||
|
// Verify it's removed from session index
|
||||||
|
artifacts := am.GetSessionArtifacts("session_1")
|
||||||
|
if len(artifacts) != 0 {
|
||||||
|
t.Error("Artifact still in session index")
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("DeleteArtifact_NotFound", func(t *testing.T) {
|
||||||
|
am := NewArtifactManager()
|
||||||
|
err := am.DeleteArtifact("nonexistent")
|
||||||
|
|
||||||
|
if err == nil {
|
||||||
|
t.Error("Expected error for nonexistent artifact")
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("DeleteSessionArtifacts", func(t *testing.T) {
|
||||||
|
am := NewArtifactManager()
|
||||||
|
am.CreateArtifact(ArtifactTypeCode, "code1", "session_1")
|
||||||
|
am.CreateArtifact(ArtifactTypeCode, "code2", "session_1")
|
||||||
|
am.CreateArtifact(ArtifactTypeCode, "code3", "session_2")
|
||||||
|
|
||||||
|
count := am.DeleteSessionArtifacts("session_1")
|
||||||
|
if count != 2 {
|
||||||
|
t.Errorf("Expected 2 artifacts deleted, got %d", count)
|
||||||
|
}
|
||||||
|
|
||||||
|
artifacts := am.GetSessionArtifacts("session_1")
|
||||||
|
if len(artifacts) != 0 {
|
||||||
|
t.Error("Session still has artifacts after deletion")
|
||||||
|
}
|
||||||
|
|
||||||
|
// session_2 should be unaffected
|
||||||
|
artifacts = am.GetSessionArtifacts("session_2")
|
||||||
|
if len(artifacts) != 1 {
|
||||||
|
t.Error("Other session was affected by deletion")
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("CountArtifacts", func(t *testing.T) {
|
||||||
|
am := NewArtifactManager()
|
||||||
|
if am.CountArtifacts() != 0 {
|
||||||
|
t.Error("Expected 0 artifacts initially")
|
||||||
|
}
|
||||||
|
|
||||||
|
am.CreateArtifact(ArtifactTypeCode, "code1", "session_1")
|
||||||
|
am.CreateArtifact(ArtifactTypeCode, "code2", "session_2")
|
||||||
|
|
||||||
|
if am.CountArtifacts() != 2 {
|
||||||
|
t.Errorf("Expected 2 artifacts, got %d", am.CountArtifacts())
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("CountSessionArtifacts", func(t *testing.T) {
|
||||||
|
am := NewArtifactManager()
|
||||||
|
am.CreateArtifact(ArtifactTypeCode, "code1", "session_1")
|
||||||
|
am.CreateArtifact(ArtifactTypeCode, "code2", "session_1")
|
||||||
|
|
||||||
|
count := am.CountSessionArtifacts("session_1")
|
||||||
|
if count != 2 {
|
||||||
|
t.Errorf("Expected 2 artifacts for session_1, got %d", count)
|
||||||
|
}
|
||||||
|
|
||||||
|
count = am.CountSessionArtifacts("nonexistent")
|
||||||
|
if count != 0 {
|
||||||
|
t.Errorf("Expected 0 artifacts for nonexistent session, got %d", count)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("GetStats", func(t *testing.T) {
|
||||||
|
am := NewArtifactManager()
|
||||||
|
am.CreateArtifact(ArtifactTypeCode, "code", "session_1")
|
||||||
|
am.CreateArtifact(ArtifactTypeImage, "image", "session_1")
|
||||||
|
am.CreateArtifact(ArtifactTypeCode, "code2", "session_2")
|
||||||
|
|
||||||
|
stats := am.GetStats()
|
||||||
|
|
||||||
|
if stats.Total != 3 {
|
||||||
|
t.Errorf("Expected 3 total artifacts, got %d", stats.Total)
|
||||||
|
}
|
||||||
|
|
||||||
|
if stats.ByType[ArtifactTypeCode] != 2 {
|
||||||
|
t.Errorf("Expected 2 code artifacts, got %d", stats.ByType[ArtifactTypeCode])
|
||||||
|
}
|
||||||
|
|
||||||
|
if stats.ByType[ArtifactTypeImage] != 1 {
|
||||||
|
t.Errorf("Expected 1 image artifact, got %d", stats.ByType[ArtifactTypeImage])
|
||||||
|
}
|
||||||
|
|
||||||
|
if stats.BySession["session_1"] != 2 {
|
||||||
|
t.Errorf("Expected 2 artifacts in session_1, got %d", stats.BySession["session_1"])
|
||||||
|
}
|
||||||
|
|
||||||
|
if stats.TotalSize == 0 {
|
||||||
|
t.Error("Expected non-zero total size")
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("ConcurrentAccess", func(t *testing.T) {
|
||||||
|
am := NewArtifactManager()
|
||||||
|
done := make(chan bool)
|
||||||
|
|
||||||
|
// Concurrent writes
|
||||||
|
for i := 0; i < 10; i++ {
|
||||||
|
go func(n int) {
|
||||||
|
am.CreateArtifact(ArtifactTypeCode, "code", "session_1")
|
||||||
|
done <- true
|
||||||
|
}(i)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Wait for all writes
|
||||||
|
for i := 0; i < 10; i++ {
|
||||||
|
<-done
|
||||||
|
}
|
||||||
|
|
||||||
|
count := am.CountSessionArtifacts("session_1")
|
||||||
|
if count != 10 {
|
||||||
|
t.Errorf("Expected 10 artifacts, got %d (concurrent writes failed)", count)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
// TestArtifactFilters tests artifact filtering
|
||||||
|
func TestArtifactFilters(t *testing.T) {
|
||||||
|
now := time.Now()
|
||||||
|
artifacts := []*Artifact{
|
||||||
|
{Type: ArtifactTypeCode, Language: "go", Size: 100, Created: now.Add(-time.Hour)},
|
||||||
|
{Type: ArtifactTypeImage, Size: 500, Created: now.Add(-time.Minute)},
|
||||||
|
{Type: ArtifactTypeCode, Language: "python", Size: 200, Created: now},
|
||||||
|
}
|
||||||
|
|
||||||
|
t.Run("FilterByType", func(t *testing.T) {
|
||||||
|
result := FilterArtifacts(artifacts, FilterArtifactsByType(ArtifactTypeCode))
|
||||||
|
if len(result) != 2 {
|
||||||
|
t.Errorf("Expected 2 code artifacts, got %d", len(result))
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("FilterByLanguage", func(t *testing.T) {
|
||||||
|
result := FilterArtifacts(artifacts, FilterArtifactsByLanguage("go"))
|
||||||
|
if len(result) != 1 {
|
||||||
|
t.Errorf("Expected 1 go artifact, got %d", len(result))
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("FilterByMinSize", func(t *testing.T) {
|
||||||
|
result := FilterArtifacts(artifacts, FilterArtifactsByMinSize(150))
|
||||||
|
if len(result) != 2 {
|
||||||
|
t.Errorf("Expected 2 artifacts >= 150 bytes, got %d", len(result))
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("FilterSince", func(t *testing.T) {
|
||||||
|
result := FilterArtifacts(artifacts, FilterArtifactsSince(now.Add(-30*time.Minute)))
|
||||||
|
if len(result) != 2 {
|
||||||
|
t.Errorf("Expected 2 recent artifacts, got %d", len(result))
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
// TestGenerateArtifactID tests ID generation
|
||||||
|
func TestGenerateArtifactID(t *testing.T) {
|
||||||
|
t.Run("DeterministicGeneration", func(t *testing.T) {
|
||||||
|
content := "test content"
|
||||||
|
id1 := generateArtifactID(content)
|
||||||
|
id2 := generateArtifactID(content)
|
||||||
|
|
||||||
|
if id1 != id2 {
|
||||||
|
t.Error("Expected same ID for same content")
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("UniqueForDifferentContent", func(t *testing.T) {
|
||||||
|
id1 := generateArtifactID("content1")
|
||||||
|
id2 := generateArtifactID("content2")
|
||||||
|
|
||||||
|
if id1 == id2 {
|
||||||
|
t.Error("Expected different IDs for different content")
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("HasPrefix", func(t *testing.T) {
|
||||||
|
id := generateArtifactID("test")
|
||||||
|
if id[:9] != "artifact_" {
|
||||||
|
t.Error("Expected artifact ID to have 'artifact_' prefix")
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
// BenchmarkArtifactManager benchmarks artifact manager operations
|
||||||
|
func BenchmarkArtifactManager(b *testing.B) {
|
||||||
|
b.Run("CreateArtifact", func(b *testing.B) {
|
||||||
|
am := NewArtifactManager()
|
||||||
|
for i := 0; i < b.N; i++ {
|
||||||
|
am.CreateArtifact(ArtifactTypeCode, "code content", "session_1")
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
b.Run("GetArtifact", func(b *testing.B) {
|
||||||
|
am := NewArtifactManager()
|
||||||
|
artifact := am.CreateArtifact(ArtifactTypeCode, "code", "session_1")
|
||||||
|
b.ResetTimer()
|
||||||
|
for i := 0; i < b.N; i++ {
|
||||||
|
_, _ = am.GetArtifact(artifact.ID)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
b.Run("GetSessionArtifacts", func(b *testing.B) {
|
||||||
|
am := NewArtifactManager()
|
||||||
|
for i := 0; i < 100; i++ {
|
||||||
|
am.CreateArtifact(ArtifactTypeCode, "code", "session_1")
|
||||||
|
}
|
||||||
|
b.ResetTimer()
|
||||||
|
for i := 0; i < b.N; i++ {
|
||||||
|
_ = am.GetSessionArtifacts("session_1")
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
b.Run("FilterArtifacts", func(b *testing.B) {
|
||||||
|
artifacts := make([]*Artifact, 100)
|
||||||
|
for i := 0; i < 100; i++ {
|
||||||
|
artifacts[i] = &Artifact{Type: ArtifactTypeCode}
|
||||||
|
}
|
||||||
|
b.ResetTimer()
|
||||||
|
for i := 0; i < b.N; i++ {
|
||||||
|
_ = FilterArtifacts(artifacts, FilterArtifactsByType(ArtifactTypeCode))
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
266
pkg/agent/interrupt.go
Normal file
266
pkg/agent/interrupt.go
Normal file
|
|
@ -0,0 +1,266 @@
|
||||||
|
// PicoClaw - Ultra-lightweight personal AI agent
|
||||||
|
// License: MIT
|
||||||
|
//
|
||||||
|
// Copyright (c) 2026 PicoClaw contributors
|
||||||
|
|
||||||
|
// DEPRECATED: This file contains the legacy interrupt handler for Phase 1.5 command-based interruption.
|
||||||
|
// The new steering architecture (nanobot-inspired) uses message injection with LLM-driven decisions.
|
||||||
|
// This code is kept for backward compatibility but will be removed in a future version.
|
||||||
|
// See: pkg/agent/interruption_checker.go for the new implementation.
|
||||||
|
|
||||||
|
package agent
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"errors"
|
||||||
|
"strings"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"github.com/sipeed/picoclaw/pkg/bus"
|
||||||
|
"github.com/sipeed/picoclaw/pkg/logger"
|
||||||
|
)
|
||||||
|
|
||||||
|
// Interrupt command constants
|
||||||
|
const (
|
||||||
|
InterruptCommandCancel = "/cancel"
|
||||||
|
InterruptCommandStop = "/stop"
|
||||||
|
InterruptCommandAbort = "/abort"
|
||||||
|
)
|
||||||
|
|
||||||
|
// ErrInterrupted is returned when agent execution is interrupted by a high-priority signal
|
||||||
|
var ErrInterrupted = errors.New("agent execution interrupted")
|
||||||
|
|
||||||
|
// InterruptSignal represents an interruption event
|
||||||
|
type InterruptSignal struct {
|
||||||
|
Type string // "user_message", "system_alert", "timeout", "cancellation"
|
||||||
|
Priority int // Priority level 1-10, higher means more urgent
|
||||||
|
Data any // Signal-specific data
|
||||||
|
Source string // Source identifier (e.g., channel:chatID)
|
||||||
|
}
|
||||||
|
|
||||||
|
// InterruptHandler checks for interruption signals during agent execution
|
||||||
|
type InterruptHandler interface {
|
||||||
|
// CheckInterruption checks if there is a pending interruption signal
|
||||||
|
// Returns nil if no interruption, or an InterruptSignal if one exists
|
||||||
|
CheckInterruption(ctx context.Context) (*InterruptSignal, error)
|
||||||
|
}
|
||||||
|
|
||||||
|
// InterruptionConfig configures interruption behavior
|
||||||
|
type InterruptionConfig struct {
|
||||||
|
Enabled bool // Whether interruption checking is enabled
|
||||||
|
CheckInterval time.Duration // Minimum interval between checks
|
||||||
|
MinPriority int // Minimum priority to trigger interruption (1-10)
|
||||||
|
MaxQueueSize int // Maximum pending messages to check
|
||||||
|
SkipInternalChan bool // Skip checking for internal channels
|
||||||
|
AllowSameSession bool // Allow interruption within the same session
|
||||||
|
MindChangeWindow time.Duration // Time window to detect "mind change" (e.g., 5 seconds)
|
||||||
|
InterruptCommands []string // Commands that trigger immediate interruption
|
||||||
|
}
|
||||||
|
|
||||||
|
// DefaultInterruptionConfig returns sensible defaults
|
||||||
|
func DefaultInterruptionConfig() InterruptionConfig {
|
||||||
|
return InterruptionConfig{
|
||||||
|
Enabled: true, // Disabled by default for backward compatibility
|
||||||
|
CheckInterval: 1 * time.Second,
|
||||||
|
MinPriority: 8, // Only interrupt for high-priority messages
|
||||||
|
MaxQueueSize: 10,
|
||||||
|
SkipInternalChan: true,
|
||||||
|
AllowSameSession: true, // Disabled by default
|
||||||
|
MindChangeWindow: 5 * time.Second,
|
||||||
|
InterruptCommands: []string{InterruptCommandCancel, InterruptCommandStop, InterruptCommandAbort},
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// BusInterruptHandler checks for new inbound messages on the message bus
|
||||||
|
type BusInterruptHandler struct {
|
||||||
|
bus *bus.MessageBus
|
||||||
|
config InterruptionConfig
|
||||||
|
currentChan string
|
||||||
|
currentChatID string
|
||||||
|
lastCheck time.Time
|
||||||
|
lastMessageTime time.Time // Track last message time for mind-change detection
|
||||||
|
}
|
||||||
|
|
||||||
|
// NewBusInterruptHandler creates a new interrupt handler that monitors the message bus
|
||||||
|
func NewBusInterruptHandler(msgBus *bus.MessageBus, config InterruptionConfig) *BusInterruptHandler {
|
||||||
|
return &BusInterruptHandler{
|
||||||
|
bus: msgBus,
|
||||||
|
config: config,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// SetContext updates the current execution context (channel and chat ID)
|
||||||
|
func (h *BusInterruptHandler) SetContext(channel, chatID string) {
|
||||||
|
h.currentChan = channel
|
||||||
|
h.currentChatID = chatID
|
||||||
|
}
|
||||||
|
|
||||||
|
// CheckInterruption implements InterruptHandler interface
|
||||||
|
func (h *BusInterruptHandler) CheckInterruption(ctx context.Context) (*InterruptSignal, error) {
|
||||||
|
// Check if interruption is enabled
|
||||||
|
if !h.config.Enabled {
|
||||||
|
return nil, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// Rate limiting: avoid checking too frequently
|
||||||
|
if time.Since(h.lastCheck) < h.config.CheckInterval {
|
||||||
|
return nil, nil
|
||||||
|
}
|
||||||
|
h.lastCheck = time.Now()
|
||||||
|
|
||||||
|
// Check context cancellation first
|
||||||
|
select {
|
||||||
|
case <-ctx.Done():
|
||||||
|
return &InterruptSignal{
|
||||||
|
Type: "cancellation",
|
||||||
|
Priority: 10, // Highest priority
|
||||||
|
Data: ctx.Err(),
|
||||||
|
Source: "context",
|
||||||
|
}, nil
|
||||||
|
default:
|
||||||
|
// Continue checking
|
||||||
|
}
|
||||||
|
|
||||||
|
// Non-blocking check for new inbound messages
|
||||||
|
// This is a simplified implementation - in production, you might want to
|
||||||
|
// peek at the bus queue without consuming messages
|
||||||
|
select {
|
||||||
|
case msg := <-h.peekInbound():
|
||||||
|
// Calculate priority based on message characteristics
|
||||||
|
priority := h.calculatePriority(msg)
|
||||||
|
|
||||||
|
// PHASE 1.5: Update lastMessageTime for mind-change detection
|
||||||
|
if msg.Channel == h.currentChan && msg.ChatID == h.currentChatID {
|
||||||
|
h.lastMessageTime = time.Now()
|
||||||
|
}
|
||||||
|
|
||||||
|
if priority >= h.config.MinPriority {
|
||||||
|
logger.InfoCF("agent", "High-priority interruption detected",
|
||||||
|
map[string]any{
|
||||||
|
"type": "user_message",
|
||||||
|
"priority": priority,
|
||||||
|
"channel": msg.Channel,
|
||||||
|
"chat_id": msg.ChatID,
|
||||||
|
})
|
||||||
|
|
||||||
|
return &InterruptSignal{
|
||||||
|
Type: "user_message",
|
||||||
|
Priority: priority,
|
||||||
|
Data: msg,
|
||||||
|
Source: msg.Channel + ":" + msg.ChatID,
|
||||||
|
}, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// Put the message back if priority is not high enough
|
||||||
|
// (In real implementation, we wouldn't consume it in the first place)
|
||||||
|
h.bus.PublishInbound(ctx, msg)
|
||||||
|
|
||||||
|
default:
|
||||||
|
// No new messages
|
||||||
|
}
|
||||||
|
|
||||||
|
return nil, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// peekInbound attempts to peek at the inbound bus without blocking
|
||||||
|
// This is a simplified implementation for demonstration
|
||||||
|
func (h *BusInterruptHandler) peekInbound() <-chan bus.InboundMessage {
|
||||||
|
// In a real implementation, you might want to add a Peek method to the MessageBus
|
||||||
|
// For now, we'll use a goroutine with timeout
|
||||||
|
ch := make(chan bus.InboundMessage, 1)
|
||||||
|
|
||||||
|
go func() {
|
||||||
|
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Millisecond)
|
||||||
|
defer cancel()
|
||||||
|
|
||||||
|
msg, ok := h.bus.ConsumeInbound(ctx)
|
||||||
|
if ok {
|
||||||
|
ch <- msg
|
||||||
|
}
|
||||||
|
close(ch)
|
||||||
|
}()
|
||||||
|
|
||||||
|
return ch
|
||||||
|
}
|
||||||
|
|
||||||
|
// calculatePriority determines the priority of a message
|
||||||
|
func (h *BusInterruptHandler) calculatePriority(msg bus.InboundMessage) int {
|
||||||
|
content := strings.ToLower(strings.TrimSpace(msg.Content))
|
||||||
|
isSameSession := (msg.Channel == h.currentChan && msg.ChatID == h.currentChatID)
|
||||||
|
|
||||||
|
// PHASE 1.5: Check for interrupt commands first (highest priority)
|
||||||
|
// Interrupt commands should ALWAYS have max priority, regardless of session
|
||||||
|
for _, cmd := range h.config.InterruptCommands {
|
||||||
|
if strings.HasPrefix(content, strings.ToLower(cmd)) {
|
||||||
|
logger.DebugCF("agent", "Interrupt command detected in priority calculation",
|
||||||
|
map[string]any{
|
||||||
|
"command": cmd,
|
||||||
|
"content": msg.Content,
|
||||||
|
"same_session": isSameSession,
|
||||||
|
"current_chan": h.currentChan,
|
||||||
|
"current_chatid": h.currentChatID,
|
||||||
|
"msg_chan": msg.Channel,
|
||||||
|
"msg_chatid": msg.ChatID,
|
||||||
|
})
|
||||||
|
return 10 // Maximum priority for explicit interrupt commands
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Default priority
|
||||||
|
priority := 5
|
||||||
|
|
||||||
|
// PHASE 1.5: Mind-change detection (boost same-session priority)
|
||||||
|
if h.config.AllowSameSession && isSameSession {
|
||||||
|
if !h.lastMessageTime.IsZero() && time.Since(h.lastMessageTime) < h.config.MindChangeWindow {
|
||||||
|
// User sent new message quickly after previous one - likely changing their mind
|
||||||
|
priority += 4 // Boost to 9, just above threshold of 8
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Higher priority for different channels (same channel/chat = lower priority)
|
||||||
|
if !isSameSession {
|
||||||
|
priority += 3
|
||||||
|
}
|
||||||
|
|
||||||
|
// Keywords that indicate urgency
|
||||||
|
urgentKeywords := []string{"urgent", "emergency", "stop", "cancel", "help"}
|
||||||
|
for _, keyword := range urgentKeywords {
|
||||||
|
if contains(msg.Content, keyword) {
|
||||||
|
priority += 2
|
||||||
|
break
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Commands typically have higher priority
|
||||||
|
if len(msg.Content) > 0 && msg.Content[0] == '/' {
|
||||||
|
priority += 1
|
||||||
|
}
|
||||||
|
|
||||||
|
// Cap at 10
|
||||||
|
if priority > 10 {
|
||||||
|
priority = 10
|
||||||
|
}
|
||||||
|
|
||||||
|
return priority
|
||||||
|
}
|
||||||
|
|
||||||
|
// contains checks if a string contains a substring (case-insensitive)
|
||||||
|
func contains(s, substr string) bool {
|
||||||
|
return len(s) >= len(substr) &&
|
||||||
|
(s == substr ||
|
||||||
|
len(s) > len(substr) &&
|
||||||
|
(s[:len(substr)] == substr || s[len(s)-len(substr):] == substr))
|
||||||
|
}
|
||||||
|
|
||||||
|
// NoOpInterruptHandler is a no-op implementation that never interrupts
|
||||||
|
type NoOpInterruptHandler struct{}
|
||||||
|
|
||||||
|
// CheckInterruption always returns nil (no interruption)
|
||||||
|
func (h *NoOpInterruptHandler) CheckInterruption(ctx context.Context) (*InterruptSignal, error) {
|
||||||
|
return nil, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// NewNoOpInterruptHandler creates a new no-op interrupt handler
|
||||||
|
func NewNoOpInterruptHandler() *NoOpInterruptHandler {
|
||||||
|
return &NoOpInterruptHandler{}
|
||||||
|
}
|
||||||
679
pkg/agent/interrupt_test.go
Normal file
679
pkg/agent/interrupt_test.go
Normal file
|
|
@ -0,0 +1,679 @@
|
||||||
|
// PicoClaw - Ultra-lightweight personal AI agent
|
||||||
|
// License: MIT
|
||||||
|
//
|
||||||
|
// Copyright (c) 2026 PicoClaw contributors
|
||||||
|
|
||||||
|
package agent
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"testing"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"github.com/sipeed/picoclaw/pkg/bus"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestInterruptSignal(t *testing.T) {
|
||||||
|
signal := InterruptSignal{
|
||||||
|
Type: "user_message",
|
||||||
|
Priority: 9,
|
||||||
|
Data: "test message",
|
||||||
|
Source: "telegram:123",
|
||||||
|
}
|
||||||
|
|
||||||
|
if signal.Type != "user_message" {
|
||||||
|
t.Errorf("Expected type 'user_message', got '%s'", signal.Type)
|
||||||
|
}
|
||||||
|
if signal.Priority != 9 {
|
||||||
|
t.Errorf("Expected priority 9, got %d", signal.Priority)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestDefaultInterruptionConfig(t *testing.T) {
|
||||||
|
config := DefaultInterruptionConfig()
|
||||||
|
|
||||||
|
// Enabled is true by default (changed from false for Phase 1.5)
|
||||||
|
if !config.Enabled {
|
||||||
|
t.Error("Expected interruption to be enabled by default")
|
||||||
|
}
|
||||||
|
if config.MinPriority != 8 {
|
||||||
|
t.Errorf("Expected min priority 8, got %d", config.MinPriority)
|
||||||
|
}
|
||||||
|
if config.CheckInterval != 1*time.Second {
|
||||||
|
t.Errorf("Expected check interval 1s, got %v", config.CheckInterval)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestNoOpInterruptHandler(t *testing.T) {
|
||||||
|
handler := NewNoOpInterruptHandler()
|
||||||
|
ctx := context.Background()
|
||||||
|
|
||||||
|
signal, err := handler.CheckInterruption(ctx)
|
||||||
|
if err != nil {
|
||||||
|
t.Errorf("Expected no error, got %v", err)
|
||||||
|
}
|
||||||
|
if signal != nil {
|
||||||
|
t.Errorf("Expected nil signal, got %+v", signal)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestBusInterruptHandler_Disabled(t *testing.T) {
|
||||||
|
msgBus := bus.NewMessageBus()
|
||||||
|
config := DefaultInterruptionConfig()
|
||||||
|
config.Enabled = false
|
||||||
|
|
||||||
|
handler := NewBusInterruptHandler(msgBus, config)
|
||||||
|
handler.SetContext("telegram", "123")
|
||||||
|
|
||||||
|
ctx := context.Background()
|
||||||
|
signal, err := handler.CheckInterruption(ctx)
|
||||||
|
if err != nil {
|
||||||
|
t.Errorf("Expected no error, got %v", err)
|
||||||
|
}
|
||||||
|
if signal != nil {
|
||||||
|
t.Errorf("Expected nil signal when disabled, got %+v", signal)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestBusInterruptHandler_ContextCancellation(t *testing.T) {
|
||||||
|
msgBus := bus.NewMessageBus()
|
||||||
|
config := DefaultInterruptionConfig()
|
||||||
|
config.Enabled = true
|
||||||
|
|
||||||
|
handler := NewBusInterruptHandler(msgBus, config)
|
||||||
|
handler.SetContext("telegram", "123")
|
||||||
|
|
||||||
|
ctx, cancel := context.WithCancel(context.Background())
|
||||||
|
cancel() // Cancel immediately
|
||||||
|
|
||||||
|
signal, err := handler.CheckInterruption(ctx)
|
||||||
|
if err != nil {
|
||||||
|
t.Errorf("Expected no error, got %v", err)
|
||||||
|
}
|
||||||
|
if signal == nil {
|
||||||
|
t.Fatal("Expected cancellation signal, got nil")
|
||||||
|
}
|
||||||
|
if signal.Type != "cancellation" {
|
||||||
|
t.Errorf("Expected type 'cancellation', got '%s'", signal.Type)
|
||||||
|
}
|
||||||
|
if signal.Priority != 10 {
|
||||||
|
t.Errorf("Expected priority 10 for cancellation, got %d", signal.Priority)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestBusInterruptHandler_RateLimiting(t *testing.T) {
|
||||||
|
msgBus := bus.NewMessageBus()
|
||||||
|
config := DefaultInterruptionConfig()
|
||||||
|
config.Enabled = true
|
||||||
|
config.CheckInterval = 1 * time.Second
|
||||||
|
|
||||||
|
handler := NewBusInterruptHandler(msgBus, config)
|
||||||
|
handler.SetContext("telegram", "123")
|
||||||
|
|
||||||
|
ctx := context.Background()
|
||||||
|
|
||||||
|
// First check
|
||||||
|
_, err := handler.CheckInterruption(ctx)
|
||||||
|
if err != nil {
|
||||||
|
t.Errorf("Expected no error, got %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Immediate second check should be skipped due to rate limiting
|
||||||
|
_, err = handler.CheckInterruption(ctx)
|
||||||
|
if err != nil {
|
||||||
|
t.Errorf("Expected no error, got %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Wait for check interval to pass
|
||||||
|
time.Sleep(1100 * time.Millisecond)
|
||||||
|
|
||||||
|
// Third check should proceed
|
||||||
|
_, err = handler.CheckInterruption(ctx)
|
||||||
|
if err != nil {
|
||||||
|
t.Errorf("Expected no error, got %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestBusInterruptHandler_SetContext(t *testing.T) {
|
||||||
|
msgBus := bus.NewMessageBus()
|
||||||
|
config := DefaultInterruptionConfig()
|
||||||
|
|
||||||
|
handler := NewBusInterruptHandler(msgBus, config)
|
||||||
|
|
||||||
|
handler.SetContext("telegram", "123")
|
||||||
|
if handler.currentChan != "telegram" {
|
||||||
|
t.Errorf("Expected channel 'telegram', got '%s'", handler.currentChan)
|
||||||
|
}
|
||||||
|
if handler.currentChatID != "123" {
|
||||||
|
t.Errorf("Expected chat ID '123', got '%s'", handler.currentChatID)
|
||||||
|
}
|
||||||
|
|
||||||
|
handler.SetContext("discord", "456")
|
||||||
|
if handler.currentChan != "discord" {
|
||||||
|
t.Errorf("Expected channel 'discord', got '%s'", handler.currentChan)
|
||||||
|
}
|
||||||
|
if handler.currentChatID != "456" {
|
||||||
|
t.Errorf("Expected chat ID '456', got '%s'", handler.currentChatID)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestCalculatePriority_DifferentChannel(t *testing.T) {
|
||||||
|
msgBus := bus.NewMessageBus()
|
||||||
|
config := DefaultInterruptionConfig()
|
||||||
|
|
||||||
|
handler := NewBusInterruptHandler(msgBus, config)
|
||||||
|
handler.SetContext("telegram", "123")
|
||||||
|
|
||||||
|
// Message from different channel should have higher priority
|
||||||
|
msg := bus.InboundMessage{
|
||||||
|
Channel: "discord",
|
||||||
|
ChatID: "456",
|
||||||
|
Content: "test message",
|
||||||
|
}
|
||||||
|
|
||||||
|
priority := handler.calculatePriority(msg)
|
||||||
|
if priority < 8 {
|
||||||
|
t.Errorf("Expected priority >= 8 for different channel, got %d", priority)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestCalculatePriority_UrgentKeywords(t *testing.T) {
|
||||||
|
msgBus := bus.NewMessageBus()
|
||||||
|
config := DefaultInterruptionConfig()
|
||||||
|
|
||||||
|
handler := NewBusInterruptHandler(msgBus, config)
|
||||||
|
handler.SetContext("telegram", "123")
|
||||||
|
|
||||||
|
testCases := []struct {
|
||||||
|
content string
|
||||||
|
keyword string
|
||||||
|
expected int // minimum expected priority
|
||||||
|
}{
|
||||||
|
{"urgent: need help", "urgent", 7},
|
||||||
|
{"emergency situation", "emergency", 7},
|
||||||
|
{"please stop", "stop", 7},
|
||||||
|
{"cancel the task", "cancel", 7},
|
||||||
|
{"help me", "help", 7},
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, tc := range testCases {
|
||||||
|
msg := bus.InboundMessage{
|
||||||
|
Channel: "telegram",
|
||||||
|
ChatID: "123",
|
||||||
|
Content: tc.content,
|
||||||
|
}
|
||||||
|
|
||||||
|
priority := handler.calculatePriority(msg)
|
||||||
|
if priority < tc.expected {
|
||||||
|
t.Errorf("Content '%s' with keyword '%s': expected priority >= %d, got %d",
|
||||||
|
tc.content, tc.keyword, tc.expected, priority)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestCalculatePriority_Command(t *testing.T) {
|
||||||
|
msgBus := bus.NewMessageBus()
|
||||||
|
config := DefaultInterruptionConfig()
|
||||||
|
|
||||||
|
handler := NewBusInterruptHandler(msgBus, config)
|
||||||
|
handler.SetContext("telegram", "123")
|
||||||
|
|
||||||
|
msg := bus.InboundMessage{
|
||||||
|
Channel: "telegram",
|
||||||
|
ChatID: "123",
|
||||||
|
Content: "/status",
|
||||||
|
}
|
||||||
|
|
||||||
|
priority := handler.calculatePriority(msg)
|
||||||
|
if priority < 6 {
|
||||||
|
t.Errorf("Expected priority >= 6 for command, got %d", priority)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestCalculatePriority_Capped(t *testing.T) {
|
||||||
|
msgBus := bus.NewMessageBus()
|
||||||
|
config := DefaultInterruptionConfig()
|
||||||
|
|
||||||
|
handler := NewBusInterruptHandler(msgBus, config)
|
||||||
|
handler.SetContext("telegram", "123")
|
||||||
|
|
||||||
|
// Message with multiple priority boosts
|
||||||
|
msg := bus.InboundMessage{
|
||||||
|
Channel: "discord", // +3
|
||||||
|
ChatID: "456",
|
||||||
|
Content: "/urgent emergency stop", // +1 (command) + 2 (urgent keyword)
|
||||||
|
}
|
||||||
|
|
||||||
|
priority := handler.calculatePriority(msg)
|
||||||
|
if priority > 10 {
|
||||||
|
t.Errorf("Expected priority capped at 10, got %d", priority)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestContains(t *testing.T) {
|
||||||
|
testCases := []struct {
|
||||||
|
s string
|
||||||
|
substr string
|
||||||
|
expected bool
|
||||||
|
}{
|
||||||
|
{"urgent message", "urgent", true},
|
||||||
|
{"this is urgent", "urgent", true},
|
||||||
|
{"URGENT", "URGENT", true},
|
||||||
|
{"emergency", "emer", true},
|
||||||
|
{"test", "testing", false},
|
||||||
|
{"", "test", false},
|
||||||
|
{"test", "", true},
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, tc := range testCases {
|
||||||
|
result := contains(tc.s, tc.substr)
|
||||||
|
if result != tc.expected {
|
||||||
|
t.Errorf("contains(%q, %q): expected %v, got %v",
|
||||||
|
tc.s, tc.substr, tc.expected, result)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestErrInterrupted(t *testing.T) {
|
||||||
|
if ErrInterrupted == nil {
|
||||||
|
t.Error("Expected ErrInterrupted to be defined")
|
||||||
|
}
|
||||||
|
if ErrInterrupted.Error() != "agent execution interrupted" {
|
||||||
|
t.Errorf("Expected error message 'agent execution interrupted', got '%s'",
|
||||||
|
ErrInterrupted.Error())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Benchmark tests
|
||||||
|
func BenchmarkCheckInterruption_NoOp(b *testing.B) {
|
||||||
|
handler := NewNoOpInterruptHandler()
|
||||||
|
ctx := context.Background()
|
||||||
|
|
||||||
|
b.ResetTimer()
|
||||||
|
for i := 0; i < b.N; i++ {
|
||||||
|
handler.CheckInterruption(ctx)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func BenchmarkCheckInterruption_Disabled(b *testing.B) {
|
||||||
|
msgBus := bus.NewMessageBus()
|
||||||
|
config := DefaultInterruptionConfig()
|
||||||
|
config.Enabled = false
|
||||||
|
|
||||||
|
handler := NewBusInterruptHandler(msgBus, config)
|
||||||
|
ctx := context.Background()
|
||||||
|
|
||||||
|
b.ResetTimer()
|
||||||
|
for i := 0; i < b.N; i++ {
|
||||||
|
handler.CheckInterruption(ctx)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func BenchmarkCheckInterruption_RateLimited(b *testing.B) {
|
||||||
|
msgBus := bus.NewMessageBus()
|
||||||
|
config := DefaultInterruptionConfig()
|
||||||
|
config.Enabled = true
|
||||||
|
config.CheckInterval = 1 * time.Hour // Ensure rate limiting kicks in
|
||||||
|
|
||||||
|
handler := NewBusInterruptHandler(msgBus, config)
|
||||||
|
handler.SetContext("telegram", "123")
|
||||||
|
ctx := context.Background()
|
||||||
|
|
||||||
|
// Prime the rate limiter
|
||||||
|
handler.CheckInterruption(ctx)
|
||||||
|
|
||||||
|
b.ResetTimer()
|
||||||
|
for i := 0; i < b.N; i++ {
|
||||||
|
handler.CheckInterruption(ctx)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func BenchmarkCalculatePriority(b *testing.B) {
|
||||||
|
msgBus := bus.NewMessageBus()
|
||||||
|
config := DefaultInterruptionConfig()
|
||||||
|
|
||||||
|
handler := NewBusInterruptHandler(msgBus, config)
|
||||||
|
handler.SetContext("telegram", "123")
|
||||||
|
|
||||||
|
msg := bus.InboundMessage{
|
||||||
|
Channel: "discord",
|
||||||
|
ChatID: "456",
|
||||||
|
Content: "urgent: please help with this emergency",
|
||||||
|
}
|
||||||
|
|
||||||
|
b.ResetTimer()
|
||||||
|
for i := 0; i < b.N; i++ {
|
||||||
|
handler.calculatePriority(msg)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ============================================================================
|
||||||
|
// Phase 1.5 Tests: Interrupt Command Support
|
||||||
|
// ============================================================================
|
||||||
|
|
||||||
|
func TestPhase15_InterruptCommands(t *testing.T) {
|
||||||
|
msgBus := bus.NewMessageBus()
|
||||||
|
config := DefaultInterruptionConfig()
|
||||||
|
config.Enabled = true
|
||||||
|
config.AllowSameSession = true
|
||||||
|
config.InterruptCommands = []string{InterruptCommandCancel, InterruptCommandStop, InterruptCommandAbort}
|
||||||
|
|
||||||
|
handler := NewBusInterruptHandler(msgBus, config)
|
||||||
|
handler.SetContext("telegram", "123")
|
||||||
|
|
||||||
|
tests := []struct {
|
||||||
|
name string
|
||||||
|
content string
|
||||||
|
expected int
|
||||||
|
}{
|
||||||
|
{"Cancel command", "/cancel", 10},
|
||||||
|
{"Stop command", "/stop", 10},
|
||||||
|
{"Abort command", "/abort", 10},
|
||||||
|
{"Cancel with uppercase", "/CANCEL", 10},
|
||||||
|
{"Cancel with text after", "/cancel please stop this", 10},
|
||||||
|
{"Stop with extra spaces", " /stop ", 10},
|
||||||
|
{"Regular message", "hello world", 5},
|
||||||
|
{"Command-like but not interrupt", "/help", 8}, // 5 + 2 (contains "help" keyword) + 1 (command) = 8
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, tt := range tests {
|
||||||
|
t.Run(tt.name, func(t *testing.T) {
|
||||||
|
msg := bus.InboundMessage{
|
||||||
|
Channel: "telegram",
|
||||||
|
ChatID: "123",
|
||||||
|
Content: tt.content,
|
||||||
|
}
|
||||||
|
priority := handler.calculatePriority(msg)
|
||||||
|
if priority != tt.expected {
|
||||||
|
t.Errorf("Expected priority %d, got %d for content '%s'", tt.expected, priority, tt.content)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestPhase15_InterruptCommandsAlwaysWork(t *testing.T) {
|
||||||
|
msgBus := bus.NewMessageBus()
|
||||||
|
config := DefaultInterruptionConfig()
|
||||||
|
config.Enabled = true
|
||||||
|
config.AllowSameSession = false // Disabled for mind-change detection
|
||||||
|
config.InterruptCommands = []string{InterruptCommandCancel}
|
||||||
|
|
||||||
|
handler := NewBusInterruptHandler(msgBus, config)
|
||||||
|
handler.SetContext("telegram", "123")
|
||||||
|
|
||||||
|
msg := bus.InboundMessage{
|
||||||
|
Channel: "telegram",
|
||||||
|
ChatID: "123",
|
||||||
|
Content: "/cancel",
|
||||||
|
}
|
||||||
|
|
||||||
|
priority := handler.calculatePriority(msg)
|
||||||
|
// Bug fix: Interrupt commands should ALWAYS work, regardless of AllowSameSession
|
||||||
|
// AllowSameSession only controls mind-change detection, not interrupt commands
|
||||||
|
if priority != 10 {
|
||||||
|
t.Errorf("Interrupt command should always work (priority 10), got %d. AllowSameSession only controls mind-change detection.", priority)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestPhase15_MindChangeDetection(t *testing.T) {
|
||||||
|
msgBus := bus.NewMessageBus()
|
||||||
|
config := DefaultInterruptionConfig()
|
||||||
|
config.Enabled = true
|
||||||
|
config.AllowSameSession = true
|
||||||
|
config.MindChangeWindow = 5 * time.Second
|
||||||
|
|
||||||
|
handler := NewBusInterruptHandler(msgBus, config)
|
||||||
|
handler.SetContext("telegram", "123")
|
||||||
|
|
||||||
|
// Simulate first message
|
||||||
|
handler.lastMessageTime = time.Now()
|
||||||
|
|
||||||
|
// Wait a moment (but within window)
|
||||||
|
time.Sleep(100 * time.Millisecond)
|
||||||
|
|
||||||
|
// Second message from same session
|
||||||
|
msg := bus.InboundMessage{
|
||||||
|
Channel: "telegram",
|
||||||
|
ChatID: "123",
|
||||||
|
Content: "actually, do something else",
|
||||||
|
}
|
||||||
|
|
||||||
|
priority := handler.calculatePriority(msg)
|
||||||
|
// Should get boosted priority (5 + 4 = 9)
|
||||||
|
if priority != 9 {
|
||||||
|
t.Errorf("Expected priority 9 for mind-change, got %d", priority)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestPhase15_MindChangeOutsideWindow(t *testing.T) {
|
||||||
|
msgBus := bus.NewMessageBus()
|
||||||
|
config := DefaultInterruptionConfig()
|
||||||
|
config.Enabled = true
|
||||||
|
config.AllowSameSession = true
|
||||||
|
config.MindChangeWindow = 100 * time.Millisecond
|
||||||
|
|
||||||
|
handler := NewBusInterruptHandler(msgBus, config)
|
||||||
|
handler.SetContext("telegram", "123")
|
||||||
|
|
||||||
|
// Simulate first message
|
||||||
|
handler.lastMessageTime = time.Now()
|
||||||
|
|
||||||
|
// Wait longer than window
|
||||||
|
time.Sleep(150 * time.Millisecond)
|
||||||
|
|
||||||
|
// Second message from same session
|
||||||
|
msg := bus.InboundMessage{
|
||||||
|
Channel: "telegram",
|
||||||
|
ChatID: "123",
|
||||||
|
Content: "another request",
|
||||||
|
}
|
||||||
|
|
||||||
|
priority := handler.calculatePriority(msg)
|
||||||
|
// Should NOT get boosted priority (still 5)
|
||||||
|
if priority != 5 {
|
||||||
|
t.Errorf("Expected priority 5 (no boost), got %d", priority)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestPhase15_MindChangeDisabled(t *testing.T) {
|
||||||
|
msgBus := bus.NewMessageBus()
|
||||||
|
config := DefaultInterruptionConfig()
|
||||||
|
config.Enabled = true
|
||||||
|
config.AllowSameSession = false // Disabled
|
||||||
|
config.MindChangeWindow = 5 * time.Second
|
||||||
|
|
||||||
|
handler := NewBusInterruptHandler(msgBus, config)
|
||||||
|
handler.SetContext("telegram", "123")
|
||||||
|
|
||||||
|
handler.lastMessageTime = time.Now()
|
||||||
|
time.Sleep(100 * time.Millisecond)
|
||||||
|
|
||||||
|
msg := bus.InboundMessage{
|
||||||
|
Channel: "telegram",
|
||||||
|
ChatID: "123",
|
||||||
|
Content: "new request",
|
||||||
|
}
|
||||||
|
|
||||||
|
priority := handler.calculatePriority(msg)
|
||||||
|
// Should not get mind-change boost when AllowSameSession is false
|
||||||
|
if priority == 9 {
|
||||||
|
t.Error("Mind-change detection should not work when AllowSameSession is false")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestPhase15_DefaultConfig(t *testing.T) {
|
||||||
|
config := DefaultInterruptionConfig()
|
||||||
|
|
||||||
|
// AllowSameSession is true by default (enables mind-change detection)
|
||||||
|
// Interrupt commands work regardless of this setting (bug fix)
|
||||||
|
if !config.AllowSameSession {
|
||||||
|
t.Error("Expected AllowSameSession to be true by default")
|
||||||
|
}
|
||||||
|
|
||||||
|
if config.MindChangeWindow != 5*time.Second {
|
||||||
|
t.Errorf("Expected MindChangeWindow to be 5s, got %v", config.MindChangeWindow)
|
||||||
|
}
|
||||||
|
|
||||||
|
expectedCommands := []string{InterruptCommandCancel, InterruptCommandStop, InterruptCommandAbort}
|
||||||
|
if len(config.InterruptCommands) != len(expectedCommands) {
|
||||||
|
t.Errorf("Expected %d interrupt commands, got %d", len(expectedCommands), len(config.InterruptCommands))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestPhase15_InterruptCommandConstants(t *testing.T) {
|
||||||
|
if InterruptCommandCancel != "/cancel" {
|
||||||
|
t.Errorf("Expected '/cancel', got '%s'", InterruptCommandCancel)
|
||||||
|
}
|
||||||
|
if InterruptCommandStop != "/stop" {
|
||||||
|
t.Errorf("Expected '/stop', got '%s'", InterruptCommandStop)
|
||||||
|
}
|
||||||
|
if InterruptCommandAbort != "/abort" {
|
||||||
|
t.Errorf("Expected '/abort', got '%s'", InterruptCommandAbort)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestPhase15_CombinedScenarios(t *testing.T) {
|
||||||
|
msgBus := bus.NewMessageBus()
|
||||||
|
config := DefaultInterruptionConfig()
|
||||||
|
config.Enabled = true
|
||||||
|
config.AllowSameSession = true
|
||||||
|
config.MindChangeWindow = 5 * time.Second
|
||||||
|
config.InterruptCommands = []string{InterruptCommandCancel, InterruptCommandStop}
|
||||||
|
|
||||||
|
handler := NewBusInterruptHandler(msgBus, config)
|
||||||
|
handler.SetContext("telegram", "123")
|
||||||
|
|
||||||
|
tests := []struct {
|
||||||
|
name string
|
||||||
|
setupFunc func()
|
||||||
|
msg bus.InboundMessage
|
||||||
|
expectedMin int
|
||||||
|
expectedMax int
|
||||||
|
description string
|
||||||
|
}{
|
||||||
|
{
|
||||||
|
name: "Interrupt command same session",
|
||||||
|
setupFunc: func() {},
|
||||||
|
msg: bus.InboundMessage{
|
||||||
|
Channel: "telegram",
|
||||||
|
ChatID: "123",
|
||||||
|
Content: "/cancel",
|
||||||
|
},
|
||||||
|
expectedMin: 10,
|
||||||
|
expectedMax: 10,
|
||||||
|
description: "Interrupt command should always get priority 10",
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "Mind change with urgent keyword",
|
||||||
|
setupFunc: func() {
|
||||||
|
handler.lastMessageTime = time.Now()
|
||||||
|
time.Sleep(100 * time.Millisecond)
|
||||||
|
},
|
||||||
|
msg: bus.InboundMessage{
|
||||||
|
Channel: "telegram",
|
||||||
|
ChatID: "123",
|
||||||
|
Content: "urgent: change of plans",
|
||||||
|
},
|
||||||
|
expectedMin: 10, // 5 + 4 (mind change) + 2 (urgent) = 11, capped at 10
|
||||||
|
expectedMax: 10,
|
||||||
|
description: "Mind change + urgent should cap at 10",
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "Different channel ignores mind change",
|
||||||
|
setupFunc: func() {
|
||||||
|
handler.lastMessageTime = time.Now()
|
||||||
|
time.Sleep(100 * time.Millisecond)
|
||||||
|
},
|
||||||
|
msg: bus.InboundMessage{
|
||||||
|
Channel: "discord",
|
||||||
|
ChatID: "456",
|
||||||
|
Content: "hello from different channel",
|
||||||
|
},
|
||||||
|
expectedMin: 8, // 5 + 3 (different channel)
|
||||||
|
expectedMax: 8,
|
||||||
|
description: "Different channel should not trigger mind-change detection",
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, tt := range tests {
|
||||||
|
t.Run(tt.name, func(t *testing.T) {
|
||||||
|
if tt.setupFunc != nil {
|
||||||
|
tt.setupFunc()
|
||||||
|
}
|
||||||
|
|
||||||
|
priority := handler.calculatePriority(tt.msg)
|
||||||
|
if priority < tt.expectedMin || priority > tt.expectedMax {
|
||||||
|
t.Errorf("%s: Expected priority between %d-%d, got %d",
|
||||||
|
tt.description, tt.expectedMin, tt.expectedMax, priority)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// TestBugFix_InterruptCommandsWorkRegardlessOfAllowSameSession tests the fix for the bug
|
||||||
|
// where /stop, /cancel, /abort commands didn't trigger CheckInterruption when AllowSameSession was false.
|
||||||
|
// Bug reported: 2026-03-02
|
||||||
|
func TestBugFix_InterruptCommandsWorkRegardlessOfAllowSameSession(t *testing.T) {
|
||||||
|
msgBus := bus.NewMessageBus()
|
||||||
|
|
||||||
|
tests := []struct {
|
||||||
|
name string
|
||||||
|
allowSameSession bool
|
||||||
|
command string
|
||||||
|
expectedPriority int
|
||||||
|
description string
|
||||||
|
}{
|
||||||
|
{
|
||||||
|
name: "Stop command with AllowSameSession=false",
|
||||||
|
allowSameSession: false,
|
||||||
|
command: "/stop",
|
||||||
|
expectedPriority: 10,
|
||||||
|
description: "Interrupt commands should work even when AllowSameSession is false",
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "Cancel command with AllowSameSession=false",
|
||||||
|
allowSameSession: false,
|
||||||
|
command: "/cancel",
|
||||||
|
expectedPriority: 10,
|
||||||
|
description: "Interrupt commands should work even when AllowSameSession is false",
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "Abort command with AllowSameSession=false",
|
||||||
|
allowSameSession: false,
|
||||||
|
command: "/abort",
|
||||||
|
expectedPriority: 10,
|
||||||
|
description: "Interrupt commands should work even when AllowSameSession is false",
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "Stop command with AllowSameSession=true",
|
||||||
|
allowSameSession: true,
|
||||||
|
command: "/stop",
|
||||||
|
expectedPriority: 10,
|
||||||
|
description: "Interrupt commands should work when AllowSameSession is true",
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, tt := range tests {
|
||||||
|
t.Run(tt.name, func(t *testing.T) {
|
||||||
|
config := DefaultInterruptionConfig()
|
||||||
|
config.Enabled = true
|
||||||
|
config.AllowSameSession = tt.allowSameSession
|
||||||
|
config.InterruptCommands = []string{InterruptCommandCancel, InterruptCommandStop, InterruptCommandAbort}
|
||||||
|
|
||||||
|
handler := NewBusInterruptHandler(msgBus, config)
|
||||||
|
handler.SetContext("telegram", "123")
|
||||||
|
|
||||||
|
msg := bus.InboundMessage{
|
||||||
|
Channel: "telegram",
|
||||||
|
ChatID: "123",
|
||||||
|
Content: tt.command,
|
||||||
|
}
|
||||||
|
|
||||||
|
priority := handler.calculatePriority(msg)
|
||||||
|
if priority != tt.expectedPriority {
|
||||||
|
t.Errorf("%s: Expected priority %d, got %d (AllowSameSession=%v)",
|
||||||
|
tt.description, tt.expectedPriority, priority, tt.allowSameSession)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
95
pkg/agent/interruption_checker.go
Normal file
95
pkg/agent/interruption_checker.go
Normal file
|
|
@ -0,0 +1,95 @@
|
||||||
|
// PicoClaw - Ultra-lightweight personal AI agent
|
||||||
|
// License: MIT
|
||||||
|
//
|
||||||
|
// Copyright (c) 2026 PicoClaw contributors
|
||||||
|
|
||||||
|
package agent
|
||||||
|
|
||||||
|
import (
|
||||||
|
"sync"
|
||||||
|
|
||||||
|
"github.com/sipeed/picoclaw/pkg/bus"
|
||||||
|
)
|
||||||
|
|
||||||
|
// InterruptionChecker manages pending interruption messages for a session.
|
||||||
|
// This is a simplified, nanobot-inspired approach that uses message injection
|
||||||
|
// instead of task cancellation.
|
||||||
|
//
|
||||||
|
// Design Philosophy:
|
||||||
|
// - Per-session queue for isolation
|
||||||
|
// - Thread-safe for concurrent access
|
||||||
|
// - Simple API: Signal, DrainAll, HasPending
|
||||||
|
// - Zero overhead when not in use
|
||||||
|
type InterruptionChecker struct {
|
||||||
|
queue []bus.InboundMessage
|
||||||
|
mu sync.Mutex
|
||||||
|
}
|
||||||
|
|
||||||
|
// NewInterruptionChecker creates a new checker for a session
|
||||||
|
func NewInterruptionChecker() *InterruptionChecker {
|
||||||
|
return &InterruptionChecker{
|
||||||
|
queue: make([]bus.InboundMessage, 0, 10), // Pre-allocate for common case
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Signal pushes a new interrupting message into the queue.
|
||||||
|
// This is called when a new message arrives for an already-active session.
|
||||||
|
func (ic *InterruptionChecker) Signal(msg bus.InboundMessage) {
|
||||||
|
ic.mu.Lock()
|
||||||
|
defer ic.mu.Unlock()
|
||||||
|
ic.queue = append(ic.queue, msg)
|
||||||
|
}
|
||||||
|
|
||||||
|
// DrainAll returns and clears all pending messages.
|
||||||
|
// This is called after tool execution to inject pending interruptions
|
||||||
|
// into the conversation.
|
||||||
|
func (ic *InterruptionChecker) DrainAll() []bus.InboundMessage {
|
||||||
|
ic.mu.Lock()
|
||||||
|
defer ic.mu.Unlock()
|
||||||
|
|
||||||
|
if len(ic.queue) == 0 {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// Copy messages to return
|
||||||
|
msgs := make([]bus.InboundMessage, len(ic.queue))
|
||||||
|
copy(msgs, ic.queue)
|
||||||
|
|
||||||
|
// Clear queue but keep capacity to avoid reallocation
|
||||||
|
ic.queue = ic.queue[:0]
|
||||||
|
|
||||||
|
return msgs
|
||||||
|
}
|
||||||
|
|
||||||
|
// HasPending returns true if there are pending interruptions
|
||||||
|
func (ic *InterruptionChecker) HasPending() bool {
|
||||||
|
ic.mu.Lock()
|
||||||
|
defer ic.mu.Unlock()
|
||||||
|
return len(ic.queue) > 0
|
||||||
|
}
|
||||||
|
|
||||||
|
// Peek returns the next message without removing it.
|
||||||
|
// Returns nil if queue is empty.
|
||||||
|
func (ic *InterruptionChecker) Peek() *bus.InboundMessage {
|
||||||
|
ic.mu.Lock()
|
||||||
|
defer ic.mu.Unlock()
|
||||||
|
|
||||||
|
if len(ic.queue) == 0 {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
return &ic.queue[0]
|
||||||
|
}
|
||||||
|
|
||||||
|
// Len returns the number of pending messages
|
||||||
|
func (ic *InterruptionChecker) Len() int {
|
||||||
|
ic.mu.Lock()
|
||||||
|
defer ic.mu.Unlock()
|
||||||
|
return len(ic.queue)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Clear removes all pending messages without returning them
|
||||||
|
func (ic *InterruptionChecker) Clear() {
|
||||||
|
ic.mu.Lock()
|
||||||
|
defer ic.mu.Unlock()
|
||||||
|
ic.queue = ic.queue[:0]
|
||||||
|
}
|
||||||
274
pkg/agent/interruption_checker_test.go
Normal file
274
pkg/agent/interruption_checker_test.go
Normal file
|
|
@ -0,0 +1,274 @@
|
||||||
|
// PicoClaw - Ultra-lightweight personal AI agent
|
||||||
|
// License: MIT
|
||||||
|
//
|
||||||
|
// Copyright (c) 2026 PicoClaw contributors
|
||||||
|
|
||||||
|
package agent
|
||||||
|
|
||||||
|
import (
|
||||||
|
"sync"
|
||||||
|
"testing"
|
||||||
|
|
||||||
|
"github.com/sipeed/picoclaw/pkg/bus"
|
||||||
|
"github.com/stretchr/testify/assert"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestNewInterruptionChecker(t *testing.T) {
|
||||||
|
checker := NewInterruptionChecker()
|
||||||
|
assert.NotNil(t, checker)
|
||||||
|
assert.False(t, checker.HasPending())
|
||||||
|
assert.Equal(t, 0, checker.Len())
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestInterruptionChecker_Signal(t *testing.T) {
|
||||||
|
checker := NewInterruptionChecker()
|
||||||
|
|
||||||
|
msg := bus.InboundMessage{
|
||||||
|
Channel: "test",
|
||||||
|
ChatID: "123",
|
||||||
|
Content: "interrupt",
|
||||||
|
}
|
||||||
|
|
||||||
|
checker.Signal(msg)
|
||||||
|
assert.True(t, checker.HasPending())
|
||||||
|
assert.Equal(t, 1, checker.Len())
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestInterruptionChecker_DrainAll(t *testing.T) {
|
||||||
|
checker := NewInterruptionChecker()
|
||||||
|
|
||||||
|
msg1 := bus.InboundMessage{Content: "msg1"}
|
||||||
|
msg2 := bus.InboundMessage{Content: "msg2"}
|
||||||
|
msg3 := bus.InboundMessage{Content: "msg3"}
|
||||||
|
|
||||||
|
checker.Signal(msg1)
|
||||||
|
checker.Signal(msg2)
|
||||||
|
checker.Signal(msg3)
|
||||||
|
|
||||||
|
assert.Equal(t, 3, checker.Len())
|
||||||
|
|
||||||
|
drained := checker.DrainAll()
|
||||||
|
assert.Len(t, drained, 3)
|
||||||
|
assert.Equal(t, "msg1", drained[0].Content)
|
||||||
|
assert.Equal(t, "msg2", drained[1].Content)
|
||||||
|
assert.Equal(t, "msg3", drained[2].Content)
|
||||||
|
|
||||||
|
// Queue should be empty after drain
|
||||||
|
assert.False(t, checker.HasPending())
|
||||||
|
assert.Equal(t, 0, checker.Len())
|
||||||
|
|
||||||
|
// Drain empty queue should return nil
|
||||||
|
drained2 := checker.DrainAll()
|
||||||
|
assert.Nil(t, drained2)
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestInterruptionChecker_Peek(t *testing.T) {
|
||||||
|
checker := NewInterruptionChecker()
|
||||||
|
|
||||||
|
// Peek empty queue
|
||||||
|
peeked := checker.Peek()
|
||||||
|
assert.Nil(t, peeked)
|
||||||
|
|
||||||
|
// Add messages
|
||||||
|
msg1 := bus.InboundMessage{Content: "first"}
|
||||||
|
msg2 := bus.InboundMessage{Content: "second"}
|
||||||
|
|
||||||
|
checker.Signal(msg1)
|
||||||
|
checker.Signal(msg2)
|
||||||
|
|
||||||
|
// Peek should return first message without removing
|
||||||
|
peeked = checker.Peek()
|
||||||
|
assert.NotNil(t, peeked)
|
||||||
|
assert.Equal(t, "first", peeked.Content)
|
||||||
|
|
||||||
|
// Queue should still have 2 messages
|
||||||
|
assert.Equal(t, 2, checker.Len())
|
||||||
|
|
||||||
|
// Peek again should return same message
|
||||||
|
peeked2 := checker.Peek()
|
||||||
|
assert.Equal(t, "first", peeked2.Content)
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestInterruptionChecker_Clear(t *testing.T) {
|
||||||
|
checker := NewInterruptionChecker()
|
||||||
|
|
||||||
|
checker.Signal(bus.InboundMessage{Content: "msg1"})
|
||||||
|
checker.Signal(bus.InboundMessage{Content: "msg2"})
|
||||||
|
|
||||||
|
assert.Equal(t, 2, checker.Len())
|
||||||
|
|
||||||
|
checker.Clear()
|
||||||
|
assert.Equal(t, 0, checker.Len())
|
||||||
|
assert.False(t, checker.HasPending())
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestInterruptionChecker_ConcurrentAccess(t *testing.T) {
|
||||||
|
checker := NewInterruptionChecker()
|
||||||
|
var wg sync.WaitGroup
|
||||||
|
|
||||||
|
// Concurrent signal
|
||||||
|
for i := 0; i < 100; i++ {
|
||||||
|
wg.Add(1)
|
||||||
|
go func(idx int) {
|
||||||
|
defer wg.Done()
|
||||||
|
checker.Signal(bus.InboundMessage{
|
||||||
|
Content: "msg",
|
||||||
|
})
|
||||||
|
}(i)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Concurrent drain
|
||||||
|
for i := 0; i < 10; i++ {
|
||||||
|
wg.Add(1)
|
||||||
|
go func() {
|
||||||
|
defer wg.Done()
|
||||||
|
checker.DrainAll()
|
||||||
|
}()
|
||||||
|
}
|
||||||
|
|
||||||
|
// Concurrent peek
|
||||||
|
for i := 0; i < 10; i++ {
|
||||||
|
wg.Add(1)
|
||||||
|
go func() {
|
||||||
|
defer wg.Done()
|
||||||
|
checker.Peek()
|
||||||
|
}()
|
||||||
|
}
|
||||||
|
|
||||||
|
// Concurrent HasPending
|
||||||
|
for i := 0; i < 10; i++ {
|
||||||
|
wg.Add(1)
|
||||||
|
go func() {
|
||||||
|
defer wg.Done()
|
||||||
|
checker.HasPending()
|
||||||
|
}()
|
||||||
|
}
|
||||||
|
|
||||||
|
wg.Wait()
|
||||||
|
|
||||||
|
// Should not panic - test passes if we get here
|
||||||
|
assert.True(t, true, "Concurrent access should not cause race conditions")
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestInterruptionChecker_PreservesOrder(t *testing.T) {
|
||||||
|
checker := NewInterruptionChecker()
|
||||||
|
|
||||||
|
// Signal messages in order
|
||||||
|
for i := 0; i < 10; i++ {
|
||||||
|
checker.Signal(bus.InboundMessage{
|
||||||
|
Content: string(rune('A' + i)), // A, B, C, ...
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
drained := checker.DrainAll()
|
||||||
|
assert.Len(t, drained, 10)
|
||||||
|
|
||||||
|
// Verify order is preserved (FIFO)
|
||||||
|
for i := 0; i < 10; i++ {
|
||||||
|
expected := string(rune('A' + i))
|
||||||
|
assert.Equal(t, expected, drained[i].Content)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestInterruptionChecker_MultipleSignalAndDrain(t *testing.T) {
|
||||||
|
checker := NewInterruptionChecker()
|
||||||
|
|
||||||
|
// First batch
|
||||||
|
checker.Signal(bus.InboundMessage{Content: "msg1"})
|
||||||
|
checker.Signal(bus.InboundMessage{Content: "msg2"})
|
||||||
|
|
||||||
|
batch1 := checker.DrainAll()
|
||||||
|
assert.Len(t, batch1, 2)
|
||||||
|
assert.False(t, checker.HasPending())
|
||||||
|
|
||||||
|
// Second batch
|
||||||
|
checker.Signal(bus.InboundMessage{Content: "msg3"})
|
||||||
|
checker.Signal(bus.InboundMessage{Content: "msg4"})
|
||||||
|
checker.Signal(bus.InboundMessage{Content: "msg5"})
|
||||||
|
|
||||||
|
batch2 := checker.DrainAll()
|
||||||
|
assert.Len(t, batch2, 3)
|
||||||
|
assert.False(t, checker.HasPending())
|
||||||
|
|
||||||
|
// Third drain should return nil
|
||||||
|
batch3 := checker.DrainAll()
|
||||||
|
assert.Nil(t, batch3)
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestInterruptionChecker_SignalAfterDrain(t *testing.T) {
|
||||||
|
checker := NewInterruptionChecker()
|
||||||
|
|
||||||
|
// Initial messages
|
||||||
|
checker.Signal(bus.InboundMessage{Content: "before1"})
|
||||||
|
checker.Signal(bus.InboundMessage{Content: "before2"})
|
||||||
|
|
||||||
|
// Drain
|
||||||
|
drained1 := checker.DrainAll()
|
||||||
|
assert.Len(t, drained1, 2)
|
||||||
|
|
||||||
|
// Signal new messages after drain
|
||||||
|
checker.Signal(bus.InboundMessage{Content: "after1"})
|
||||||
|
checker.Signal(bus.InboundMessage{Content: "after2"})
|
||||||
|
|
||||||
|
// Should have new messages
|
||||||
|
assert.True(t, checker.HasPending())
|
||||||
|
assert.Equal(t, 2, checker.Len())
|
||||||
|
|
||||||
|
drained2 := checker.DrainAll()
|
||||||
|
assert.Len(t, drained2, 2)
|
||||||
|
assert.Equal(t, "after1", drained2[0].Content)
|
||||||
|
assert.Equal(t, "after2", drained2[1].Content)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Benchmark tests
|
||||||
|
func BenchmarkInterruptionChecker_Signal(b *testing.B) {
|
||||||
|
checker := NewInterruptionChecker()
|
||||||
|
msg := bus.InboundMessage{Content: "test"}
|
||||||
|
|
||||||
|
b.ResetTimer()
|
||||||
|
for i := 0; i < b.N; i++ {
|
||||||
|
checker.Signal(msg)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func BenchmarkInterruptionChecker_DrainAll(b *testing.B) {
|
||||||
|
checker := NewInterruptionChecker()
|
||||||
|
|
||||||
|
// Pre-populate
|
||||||
|
for i := 0; i < 10; i++ {
|
||||||
|
checker.Signal(bus.InboundMessage{Content: "msg"})
|
||||||
|
}
|
||||||
|
|
||||||
|
b.ResetTimer()
|
||||||
|
for i := 0; i < b.N; i++ {
|
||||||
|
checker.DrainAll()
|
||||||
|
// Repopulate for next iteration
|
||||||
|
for j := 0; j < 10; j++ {
|
||||||
|
checker.Signal(bus.InboundMessage{Content: "msg"})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func BenchmarkInterruptionChecker_HasPending(b *testing.B) {
|
||||||
|
checker := NewInterruptionChecker()
|
||||||
|
checker.Signal(bus.InboundMessage{Content: "test"})
|
||||||
|
|
||||||
|
b.ResetTimer()
|
||||||
|
for i := 0; i < b.N; i++ {
|
||||||
|
checker.HasPending()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func BenchmarkInterruptionChecker_Concurrent(b *testing.B) {
|
||||||
|
checker := NewInterruptionChecker()
|
||||||
|
|
||||||
|
b.RunParallel(func(pb *testing.PB) {
|
||||||
|
for pb.Next() {
|
||||||
|
checker.Signal(bus.InboundMessage{Content: "test"})
|
||||||
|
checker.HasPending()
|
||||||
|
if checker.Len() > 100 {
|
||||||
|
checker.DrainAll()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
@ -42,6 +42,15 @@ type AgentLoop struct {
|
||||||
fallback *providers.FallbackChain
|
fallback *providers.FallbackChain
|
||||||
channelManager *channels.Manager
|
channelManager *channels.Manager
|
||||||
mediaStore media.MediaStore
|
mediaStore media.MediaStore
|
||||||
|
|
||||||
|
// Legacy interrupt handling (to be deprecated)
|
||||||
|
interruptHandler InterruptHandler // Interrupt handler for dynamic task management
|
||||||
|
taskManager *TaskManager // Task manager for concurrent task tracking (Phase 2)
|
||||||
|
|
||||||
|
// New steering architecture (nanobot-inspired)
|
||||||
|
enableSteering bool // Opt-in flag for steering feature
|
||||||
|
interruptCheckers map[string]*InterruptionChecker // Per-session interrupt queues
|
||||||
|
checkersMu sync.RWMutex // Protects interruptCheckers map
|
||||||
}
|
}
|
||||||
|
|
||||||
// processOptions configures how a message is processed
|
// processOptions configures how a message is processed
|
||||||
|
|
@ -75,14 +84,33 @@ func NewAgentLoop(cfg *config.Config, msgBus *bus.MessageBus, provider providers
|
||||||
stateManager = state.NewManager(defaultAgent.Workspace)
|
stateManager = state.NewManager(defaultAgent.Workspace)
|
||||||
}
|
}
|
||||||
|
|
||||||
return &AgentLoop{
|
al := &AgentLoop{
|
||||||
bus: msgBus,
|
bus: msgBus,
|
||||||
cfg: cfg,
|
cfg: cfg,
|
||||||
registry: registry,
|
registry: registry,
|
||||||
state: stateManager,
|
state: stateManager,
|
||||||
summarizing: sync.Map{},
|
summarizing: sync.Map{},
|
||||||
fallback: fallbackChain,
|
fallback: fallbackChain,
|
||||||
|
|
||||||
|
// New steering architecture (nanobot-inspired)
|
||||||
|
enableSteering: cfg.Agents.Defaults.EnableSteering,
|
||||||
|
interruptCheckers: make(map[string]*InterruptionChecker),
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// Legacy components (DEPRECATED - only initialized when new steering is disabled)
|
||||||
|
if !cfg.Agents.Defaults.EnableSteering {
|
||||||
|
// Legacy interrupt handler (Phase 1.5)
|
||||||
|
al.interruptHandler = NewBusInterruptHandler(msgBus, DefaultInterruptionConfig())
|
||||||
|
|
||||||
|
// Legacy TaskManager (Phase 2)
|
||||||
|
maxConcurrent := cfg.Agents.Defaults.MaxConcurrentTasks
|
||||||
|
if maxConcurrent < 0 {
|
||||||
|
maxConcurrent = 0 // Ensure no negative values, 0 = unlimited
|
||||||
|
}
|
||||||
|
al.taskManager = NewTaskManager(maxConcurrent)
|
||||||
|
}
|
||||||
|
|
||||||
|
return al
|
||||||
}
|
}
|
||||||
|
|
||||||
// registerSharedTools registers tools that are shared across all agents (web, message, spawn).
|
// registerSharedTools registers tools that are shared across all agents (web, message, spawn).
|
||||||
|
|
@ -170,9 +198,24 @@ func registerSharedTools(
|
||||||
func (al *AgentLoop) Run(ctx context.Context) error {
|
func (al *AgentLoop) Run(ctx context.Context) error {
|
||||||
al.running.Store(true)
|
al.running.Store(true)
|
||||||
|
|
||||||
|
// Legacy: Start task cleanup and steering loop only when new steering is disabled
|
||||||
|
if !al.enableSteering {
|
||||||
|
// Phase 2: Start task cleanup goroutine
|
||||||
|
go al.runTaskCleanup(ctx)
|
||||||
|
|
||||||
|
// Phase 2: Start steering loop if enabled
|
||||||
|
if al.cfg.Agents.Defaults.EnableSteeringLoop {
|
||||||
|
go al.runSteeringLoop(ctx)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
for al.running.Load() {
|
for al.running.Load() {
|
||||||
select {
|
select {
|
||||||
case <-ctx.Done():
|
case <-ctx.Done():
|
||||||
|
// Legacy: Wait for running tasks to complete (only if using TaskManager)
|
||||||
|
if !al.enableSteering && al.taskManager != nil {
|
||||||
|
al.waitForRunningTasks(5 * time.Second)
|
||||||
|
}
|
||||||
return nil
|
return nil
|
||||||
default:
|
default:
|
||||||
msg, ok := al.bus.ConsumeInbound(ctx)
|
msg, ok := al.bus.ConsumeInbound(ctx)
|
||||||
|
|
@ -180,8 +223,37 @@ func (al *AgentLoop) Run(ctx context.Context) error {
|
||||||
continue
|
continue
|
||||||
}
|
}
|
||||||
|
|
||||||
// Process message
|
// ===== NEW: Steering Architecture - Check for active session =====
|
||||||
func() {
|
if al.enableSteering {
|
||||||
|
// Get session key for this message (needs routing resolution)
|
||||||
|
sessionKey := al.getSessionKeyForMessage(msg)
|
||||||
|
|
||||||
|
// Check if this session has an active checker (task is running)
|
||||||
|
if al.hasActiveChecker(sessionKey) {
|
||||||
|
// Session is active, signal interruption instead of creating new task
|
||||||
|
checker := al.getOrCreateChecker(sessionKey)
|
||||||
|
checker.Signal(msg)
|
||||||
|
|
||||||
|
logger.InfoCF("agent", "Steering: signaled interruption for active session",
|
||||||
|
map[string]any{
|
||||||
|
"session_key": sessionKey,
|
||||||
|
"channel": msg.Channel,
|
||||||
|
"chat_id": msg.ChatID,
|
||||||
|
"content_preview": utils.Truncate(msg.Content, 60),
|
||||||
|
})
|
||||||
|
continue // Don't process as new message
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Phase 2: Process message asynchronously
|
||||||
|
go func(msg bus.InboundMessage) {
|
||||||
|
defer func() {
|
||||||
|
if r := recover(); r != nil {
|
||||||
|
logger.ErrorCF("agent", "Panic in message processing",
|
||||||
|
map[string]any{"error": r, "channel": msg.Channel, "chat_id": msg.ChatID})
|
||||||
|
}
|
||||||
|
}()
|
||||||
|
|
||||||
// TODO: Re-enable media cleanup after inbound media is properly consumed by the agent.
|
// TODO: Re-enable media cleanup after inbound media is properly consumed by the agent.
|
||||||
// Currently disabled because files are deleted before the LLM can access their content.
|
// Currently disabled because files are deleted before the LLM can access their content.
|
||||||
// defer func() {
|
// defer func() {
|
||||||
|
|
@ -197,6 +269,16 @@ func (al *AgentLoop) Run(ctx context.Context) error {
|
||||||
|
|
||||||
response, err := al.processMessage(ctx, msg)
|
response, err := al.processMessage(ctx, msg)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
|
// Phase 2: Don't send error message if task was canceled
|
||||||
|
// (user already received confirmation from /stop command)
|
||||||
|
if errors.Is(err, context.Canceled) {
|
||||||
|
logger.InfoCF("agent", "Task canceled, skipping error response",
|
||||||
|
map[string]any{
|
||||||
|
"channel": msg.Channel,
|
||||||
|
"chat_id": msg.ChatID,
|
||||||
|
})
|
||||||
|
return // Silent return on cancellation
|
||||||
|
}
|
||||||
response = fmt.Sprintf("Error processing message: %v", err)
|
response = fmt.Sprintf("Error processing message: %v", err)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
@ -234,7 +316,7 @@ func (al *AgentLoop) Run(ctx context.Context) error {
|
||||||
)
|
)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}()
|
}(msg)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
@ -245,6 +327,146 @@ func (al *AgentLoop) Stop() {
|
||||||
al.running.Store(false)
|
al.running.Store(false)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// runTaskCleanup periodically cleans up old completed tasks (Phase 2)
|
||||||
|
func (al *AgentLoop) runTaskCleanup(ctx context.Context) {
|
||||||
|
// Get cleanup interval from config, default to 5 minutes
|
||||||
|
intervalMins := al.cfg.Agents.Defaults.TaskCleanupIntervalMins
|
||||||
|
if intervalMins <= 0 {
|
||||||
|
intervalMins = 5
|
||||||
|
}
|
||||||
|
|
||||||
|
// Get retention time from config, default to 1 hour
|
||||||
|
retentionHours := al.cfg.Agents.Defaults.TaskRetentionHours
|
||||||
|
if retentionHours <= 0 {
|
||||||
|
retentionHours = 1
|
||||||
|
}
|
||||||
|
|
||||||
|
ticker := time.NewTicker(time.Duration(intervalMins) * time.Minute)
|
||||||
|
defer ticker.Stop()
|
||||||
|
|
||||||
|
logger.InfoCF("agent", "Task cleanup loop started",
|
||||||
|
map[string]any{
|
||||||
|
"interval_mins": intervalMins,
|
||||||
|
"retention_hours": retentionHours,
|
||||||
|
})
|
||||||
|
|
||||||
|
for {
|
||||||
|
select {
|
||||||
|
case <-ctx.Done():
|
||||||
|
logger.InfoCF("agent", "Task cleanup loop stopped", nil)
|
||||||
|
return
|
||||||
|
case <-ticker.C:
|
||||||
|
removed := al.taskManager.Cleanup(time.Duration(retentionHours) * time.Hour)
|
||||||
|
if removed > 0 {
|
||||||
|
logger.DebugCF("agent", "Cleaned up old tasks",
|
||||||
|
map[string]any{"removed": removed})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// waitForRunningTasks waits for all running tasks to complete or timeout (Phase 2)
|
||||||
|
func (al *AgentLoop) waitForRunningTasks(timeout time.Duration) {
|
||||||
|
deadline := time.Now().Add(timeout)
|
||||||
|
for {
|
||||||
|
tasks := al.taskManager.GetRunningTasks()
|
||||||
|
if len(tasks) == 0 {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
if time.Now().After(deadline) {
|
||||||
|
logger.WarnCF("agent", "Timeout waiting for tasks, some may be abandoned",
|
||||||
|
map[string]any{"running_tasks": len(tasks)})
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
time.Sleep(100 * time.Millisecond)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// runSteeringLoop monitors for interrupt signals and cancels tasks (Phase 2 Step 5)
|
||||||
|
func (al *AgentLoop) runSteeringLoop(ctx context.Context) {
|
||||||
|
// Get interval from config, default to 500ms
|
||||||
|
intervalMs := al.cfg.Agents.Defaults.SteeringLoopIntervalMs
|
||||||
|
if intervalMs <= 0 {
|
||||||
|
intervalMs = 500
|
||||||
|
}
|
||||||
|
|
||||||
|
ticker := time.NewTicker(time.Duration(intervalMs) * time.Millisecond)
|
||||||
|
defer ticker.Stop()
|
||||||
|
|
||||||
|
logger.InfoCF("agent", "Steering loop started",
|
||||||
|
map[string]any{"interval_ms": intervalMs})
|
||||||
|
|
||||||
|
for {
|
||||||
|
select {
|
||||||
|
case <-ctx.Done():
|
||||||
|
logger.InfoCF("agent", "Steering loop stopped", nil)
|
||||||
|
return
|
||||||
|
case <-ticker.C:
|
||||||
|
al.checkAndHandleInterrupts(ctx)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// checkAndHandleInterrupts checks for interrupt signals and handles them (Phase 2 Step 5)
|
||||||
|
func (al *AgentLoop) checkAndHandleInterrupts(ctx context.Context) {
|
||||||
|
if al.interruptHandler == nil {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
// Check for interruption signal
|
||||||
|
signal, err := al.interruptHandler.CheckInterruption(ctx)
|
||||||
|
if err != nil {
|
||||||
|
logger.WarnCF("agent", "Interrupt check failed",
|
||||||
|
map[string]any{"error": err.Error()})
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
if signal == nil || signal.Priority < 8 {
|
||||||
|
return // No interrupt or priority not high enough
|
||||||
|
}
|
||||||
|
|
||||||
|
// High-priority interrupt detected
|
||||||
|
logger.InfoCF("agent", "High-priority interrupt detected by steering loop",
|
||||||
|
map[string]any{
|
||||||
|
"type": signal.Type,
|
||||||
|
"priority": signal.Priority,
|
||||||
|
"source": signal.Source,
|
||||||
|
})
|
||||||
|
|
||||||
|
// If it's a user message interrupt, check if we need to cancel existing tasks
|
||||||
|
if signal.Type == "user_message" {
|
||||||
|
if msg, ok := signal.Data.(bus.InboundMessage); ok {
|
||||||
|
// Cancel all running tasks for the same session
|
||||||
|
canceled := al.taskManager.CancelAllTasksForSession(msg.Channel, msg.ChatID)
|
||||||
|
|
||||||
|
if canceled > 0 {
|
||||||
|
logger.InfoCF("agent", "Canceled tasks for new high-priority message",
|
||||||
|
map[string]any{
|
||||||
|
"channel": msg.Channel,
|
||||||
|
"chat_id": msg.ChatID,
|
||||||
|
"canceled_count": canceled,
|
||||||
|
})
|
||||||
|
|
||||||
|
// Send notification to user
|
||||||
|
al.bus.PublishOutbound(ctx, bus.OutboundMessage{
|
||||||
|
Channel: msg.Channel,
|
||||||
|
ChatID: msg.ChatID,
|
||||||
|
Content: "Previous task interrupted. Processing your new request...",
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
// Put the interrupt message back in the queue for normal processing
|
||||||
|
go func() {
|
||||||
|
pubCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||||
|
defer cancel()
|
||||||
|
al.bus.PublishInbound(pubCtx, msg)
|
||||||
|
}()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
func (al *AgentLoop) RegisterTool(tool tools.Tool) {
|
func (al *AgentLoop) RegisterTool(tool tools.Tool) {
|
||||||
for _, agentID := range al.registry.ListAgentIDs() {
|
for _, agentID := range al.registry.ListAgentIDs() {
|
||||||
if agent, ok := al.registry.GetAgent(agentID); ok {
|
if agent, ok := al.registry.GetAgent(agentID); ok {
|
||||||
|
|
@ -253,6 +475,104 @@ func (al *AgentLoop) RegisterTool(tool tools.Tool) {
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// ===== Steering Architecture: InterruptionChecker Management =====
|
||||||
|
|
||||||
|
// getOrCreateChecker gets or creates an interruption checker for a session.
|
||||||
|
// Thread-safe with double-checked locking pattern.
|
||||||
|
func (al *AgentLoop) getOrCreateChecker(sessionKey string) *InterruptionChecker {
|
||||||
|
// Fast path: read lock
|
||||||
|
al.checkersMu.RLock()
|
||||||
|
checker, exists := al.interruptCheckers[sessionKey]
|
||||||
|
al.checkersMu.RUnlock()
|
||||||
|
|
||||||
|
if exists {
|
||||||
|
return checker
|
||||||
|
}
|
||||||
|
|
||||||
|
// Slow path: write lock
|
||||||
|
al.checkersMu.Lock()
|
||||||
|
defer al.checkersMu.Unlock()
|
||||||
|
|
||||||
|
// Double-check after acquiring write lock
|
||||||
|
if checker, exists := al.interruptCheckers[sessionKey]; exists {
|
||||||
|
return checker
|
||||||
|
}
|
||||||
|
|
||||||
|
// Create new checker
|
||||||
|
checker = NewInterruptionChecker()
|
||||||
|
al.interruptCheckers[sessionKey] = checker
|
||||||
|
|
||||||
|
logger.DebugCF("agent", "Created interruption checker for session",
|
||||||
|
map[string]any{"session_key": sessionKey})
|
||||||
|
|
||||||
|
return checker
|
||||||
|
}
|
||||||
|
|
||||||
|
// formatInterruptionInjection formats pending interruption messages for injection into conversation.
|
||||||
|
// This follows nanobot's pattern of providing context to the LLM about the interruption.
|
||||||
|
func formatInterruptionInjection(pending []bus.InboundMessage) string {
|
||||||
|
if len(pending) == 0 {
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
|
||||||
|
var combined strings.Builder
|
||||||
|
for i, msg := range pending {
|
||||||
|
if i > 0 {
|
||||||
|
combined.WriteString("\n\n---\n\n")
|
||||||
|
}
|
||||||
|
combined.WriteString(msg.Content)
|
||||||
|
}
|
||||||
|
|
||||||
|
injection := "[The user just sent a new message while you were working. " +
|
||||||
|
"Read it and decide: continue current work, switch to the new request, or address both.]\n\n" +
|
||||||
|
combined.String()
|
||||||
|
|
||||||
|
return injection
|
||||||
|
}
|
||||||
|
|
||||||
|
// removeChecker removes a checker when session completes.
|
||||||
|
// This prevents memory leaks for long-running processes.
|
||||||
|
func (al *AgentLoop) removeChecker(sessionKey string) {
|
||||||
|
al.checkersMu.Lock()
|
||||||
|
defer al.checkersMu.Unlock()
|
||||||
|
|
||||||
|
if _, exists := al.interruptCheckers[sessionKey]; exists {
|
||||||
|
delete(al.interruptCheckers, sessionKey)
|
||||||
|
logger.DebugCF("agent", "Removed interruption checker for session",
|
||||||
|
map[string]any{"session_key": sessionKey})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// hasActiveChecker checks if a session has an active interruption checker.
|
||||||
|
// This indicates the session is currently processing a message.
|
||||||
|
func (al *AgentLoop) hasActiveChecker(sessionKey string) bool {
|
||||||
|
al.checkersMu.RLock()
|
||||||
|
defer al.checkersMu.RUnlock()
|
||||||
|
_, exists := al.interruptCheckers[sessionKey]
|
||||||
|
return exists
|
||||||
|
}
|
||||||
|
|
||||||
|
// getSessionKeyForMessage resolves the session key for a message using routing logic.
|
||||||
|
// This is needed to check if the session has an active task before creating a new one.
|
||||||
|
func (al *AgentLoop) getSessionKeyForMessage(msg bus.InboundMessage) string {
|
||||||
|
// If message already has an agent-scoped session key, use it
|
||||||
|
if msg.SessionKey != "" && strings.HasPrefix(msg.SessionKey, "agent:") {
|
||||||
|
return msg.SessionKey
|
||||||
|
}
|
||||||
|
|
||||||
|
// Otherwise, resolve via routing
|
||||||
|
route := al.registry.ResolveRoute(routing.RouteInput{
|
||||||
|
Channel: msg.Channel,
|
||||||
|
AccountID: msg.Metadata["account_id"],
|
||||||
|
Peer: extractPeer(msg),
|
||||||
|
ParentPeer: extractParentPeer(msg),
|
||||||
|
GuildID: msg.Metadata["guild_id"],
|
||||||
|
TeamID: msg.Metadata["team_id"],
|
||||||
|
})
|
||||||
|
|
||||||
|
return route.SessionKey
|
||||||
|
}
|
||||||
|
|
||||||
func (al *AgentLoop) SetChannelManager(cm *channels.Manager) {
|
func (al *AgentLoop) SetChannelManager(cm *channels.Manager) {
|
||||||
al.channelManager = cm
|
al.channelManager = cm
|
||||||
}
|
}
|
||||||
|
|
@ -262,6 +582,12 @@ func (al *AgentLoop) SetMediaStore(s media.MediaStore) {
|
||||||
al.mediaStore = s
|
al.mediaStore = s
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// SetInterruptHandler sets or replaces the interrupt handler for the agent loop.
|
||||||
|
// This allows dynamic configuration of interruption behavior at runtime.
|
||||||
|
func (al *AgentLoop) SetInterruptHandler(handler InterruptHandler) {
|
||||||
|
al.interruptHandler = handler
|
||||||
|
}
|
||||||
|
|
||||||
// inferMediaType determines the media type ("image", "audio", "video", "file")
|
// inferMediaType determines the media type ("image", "audio", "video", "file")
|
||||||
// from a filename and MIME content type.
|
// from a filename and MIME content type.
|
||||||
func inferMediaType(filename, contentType string) string {
|
func inferMediaType(filename, contentType string) string {
|
||||||
|
|
@ -364,6 +690,54 @@ func (al *AgentLoop) processMessage(ctx context.Context, msg bus.InboundMessage)
|
||||||
"session_key": msg.SessionKey,
|
"session_key": msg.SessionKey,
|
||||||
})
|
})
|
||||||
|
|
||||||
|
// New steering architecture: Direct processing without task management
|
||||||
|
if al.enableSteering {
|
||||||
|
return al.processMessageDirect(ctx, msg)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Legacy: Phase 2 task management
|
||||||
|
priority := 5 // Default priority
|
||||||
|
if busHandler, ok := al.interruptHandler.(*BusInterruptHandler); ok {
|
||||||
|
priority = busHandler.calculatePriority(msg)
|
||||||
|
}
|
||||||
|
|
||||||
|
task := NewTask(msg, priority)
|
||||||
|
if err := al.taskManager.AddTask(task); err != nil {
|
||||||
|
return "", fmt.Errorf("failed to add task: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Start task (synchronous mode for now - Phase 2 Step 2)
|
||||||
|
if err := al.taskManager.StartTask(task.ID, ctx); err != nil {
|
||||||
|
return "", fmt.Errorf("failed to start task: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Use task's context for cancellation support
|
||||||
|
taskCtx := task.Context()
|
||||||
|
|
||||||
|
// Defer task completion/failure to ensure it's always updated
|
||||||
|
defer func() {
|
||||||
|
// Check if task is still running (not already completed/failed/canceled)
|
||||||
|
if taskObj, exists := al.taskManager.GetTask(task.ID); exists {
|
||||||
|
if taskObj.Status == TaskStatusRunning {
|
||||||
|
al.taskManager.CompleteTask(task.ID)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}()
|
||||||
|
|
||||||
|
// Delegate to helper function for actual processing
|
||||||
|
response, err := al.processMessageWithTask(taskCtx, task, msg)
|
||||||
|
// Update task status based on result
|
||||||
|
if err != nil {
|
||||||
|
al.taskManager.FailTask(task.ID, err)
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
|
||||||
|
// Task will be completed by defer
|
||||||
|
return response, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// processMessageDirect handles message processing for new steering architecture (no task management)
|
||||||
|
func (al *AgentLoop) processMessageDirect(ctx context.Context, msg bus.InboundMessage) (string, error) {
|
||||||
// Route system messages to processSystemMessage
|
// Route system messages to processSystemMessage
|
||||||
if msg.Channel == "system" {
|
if msg.Channel == "system" {
|
||||||
return al.processSystemMessage(ctx, msg)
|
return al.processSystemMessage(ctx, msg)
|
||||||
|
|
@ -392,6 +766,61 @@ func (al *AgentLoop) processMessage(ctx context.Context, msg bus.InboundMessage)
|
||||||
return "", fmt.Errorf("no agent available for route (agent_id=%s)", route.AgentID)
|
return "", fmt.Errorf("no agent available for route (agent_id=%s)", route.AgentID)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// Reset message-tool state for this round
|
||||||
|
if tool, ok := agent.Tools.Get("message"); ok {
|
||||||
|
if mt, ok := tool.(tools.ContextualTool); ok {
|
||||||
|
mt.SetContext(msg.Channel, msg.ChatID)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Use routed session key, but honor pre-set agent-scoped keys
|
||||||
|
sessionKey := route.SessionKey
|
||||||
|
if msg.SessionKey != "" && strings.HasPrefix(msg.SessionKey, "agent:") {
|
||||||
|
sessionKey = msg.SessionKey
|
||||||
|
}
|
||||||
|
|
||||||
|
return al.runAgentLoop(ctx, agent, processOptions{
|
||||||
|
SessionKey: sessionKey,
|
||||||
|
Channel: msg.Channel,
|
||||||
|
ChatID: msg.ChatID,
|
||||||
|
UserMessage: msg.Content,
|
||||||
|
EnableSummary: true,
|
||||||
|
SendResponse: false,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
// processMessageWithTask handles the actual message processing logic with task context (LEGACY)
|
||||||
|
func (al *AgentLoop) processMessageWithTask(ctx context.Context, task *Task, msg bus.InboundMessage) (string, error) {
|
||||||
|
// Route system messages to processSystemMessage
|
||||||
|
if msg.Channel == "system" {
|
||||||
|
return al.processSystemMessage(ctx, msg)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Check for commands
|
||||||
|
if response, handled := al.handleCommand(ctx, msg); handled {
|
||||||
|
// Mark task as a command task (for filtering in cancellation counts)
|
||||||
|
task.Metadata["is_command"] = true
|
||||||
|
return response, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// Route to determine agent and session key
|
||||||
|
route := al.registry.ResolveRoute(routing.RouteInput{
|
||||||
|
Channel: msg.Channel,
|
||||||
|
AccountID: msg.Metadata["account_id"],
|
||||||
|
Peer: extractPeer(msg),
|
||||||
|
ParentPeer: extractParentPeer(msg),
|
||||||
|
GuildID: msg.Metadata["guild_id"],
|
||||||
|
TeamID: msg.Metadata["team_id"],
|
||||||
|
})
|
||||||
|
|
||||||
|
agent, ok := al.registry.GetAgent(route.AgentID)
|
||||||
|
if !ok {
|
||||||
|
agent = al.registry.GetDefaultAgent()
|
||||||
|
}
|
||||||
|
if agent == nil {
|
||||||
|
return "", fmt.Errorf("no agent available for route (agent_id=%s)", route.AgentID)
|
||||||
|
}
|
||||||
|
|
||||||
// Reset message-tool state for this round so we don't skip publishing due to a previous round.
|
// Reset message-tool state for this round so we don't skip publishing due to a previous round.
|
||||||
if tool, ok := agent.Tools.Get("message"); ok {
|
if tool, ok := agent.Tools.Get("message"); ok {
|
||||||
if mt, ok := tool.(tools.ContextualTool); ok {
|
if mt, ok := tool.(tools.ContextualTool); ok {
|
||||||
|
|
@ -410,6 +839,7 @@ func (al *AgentLoop) processMessage(ctx context.Context, msg bus.InboundMessage)
|
||||||
"agent_id": agent.ID,
|
"agent_id": agent.ID,
|
||||||
"session_key": sessionKey,
|
"session_key": sessionKey,
|
||||||
"matched_by": route.MatchedBy,
|
"matched_by": route.MatchedBy,
|
||||||
|
"task_id": task.ID,
|
||||||
})
|
})
|
||||||
|
|
||||||
return al.runAgentLoop(ctx, agent, processOptions{
|
return al.runAgentLoop(ctx, agent, processOptions{
|
||||||
|
|
@ -484,7 +914,28 @@ func (al *AgentLoop) processSystemMessage(ctx context.Context, msg bus.InboundMe
|
||||||
|
|
||||||
// runAgentLoop is the core message processing logic.
|
// runAgentLoop is the core message processing logic.
|
||||||
func (al *AgentLoop) runAgentLoop(ctx context.Context, agent *AgentInstance, opts processOptions) (string, error) {
|
func (al *AgentLoop) runAgentLoop(ctx context.Context, agent *AgentInstance, opts processOptions) (string, error) {
|
||||||
// 0. Record last channel for heartbeat notifications (skip internal channels)
|
// Phase 2 Step 4: Check if context is already canceled before starting
|
||||||
|
select {
|
||||||
|
case <-ctx.Done():
|
||||||
|
return "", fmt.Errorf("task canceled before execution: %w", ctx.Err())
|
||||||
|
default:
|
||||||
|
}
|
||||||
|
|
||||||
|
// ===== NEW: Steering Architecture - Setup checker for this session =====
|
||||||
|
if al.enableSteering {
|
||||||
|
// Create checker to signal this session is active
|
||||||
|
al.getOrCreateChecker(opts.SessionKey)
|
||||||
|
|
||||||
|
// Cleanup checker when done
|
||||||
|
defer al.removeChecker(opts.SessionKey)
|
||||||
|
}
|
||||||
|
|
||||||
|
// 0a. Update interrupt handler context
|
||||||
|
if busHandler, ok := al.interruptHandler.(*BusInterruptHandler); ok {
|
||||||
|
busHandler.SetContext(opts.Channel, opts.ChatID)
|
||||||
|
}
|
||||||
|
|
||||||
|
// 0b. Record last channel for heartbeat notifications (skip internal channels)
|
||||||
if opts.Channel != "" && opts.ChatID != "" {
|
if opts.Channel != "" && opts.ChatID != "" {
|
||||||
// Don't record internal channels (cli, system, subagent)
|
// Don't record internal channels (cli, system, subagent)
|
||||||
if !constants.IsInternalChannel(opts.Channel) {
|
if !constants.IsInternalChannel(opts.Channel) {
|
||||||
|
|
@ -504,6 +955,22 @@ func (al *AgentLoop) runAgentLoop(ctx context.Context, agent *AgentInstance, opt
|
||||||
if !opts.NoHistory {
|
if !opts.NoHistory {
|
||||||
history = agent.Sessions.GetHistory(opts.SessionKey)
|
history = agent.Sessions.GetHistory(opts.SessionKey)
|
||||||
summary = agent.Sessions.GetSummary(opts.SessionKey)
|
summary = agent.Sessions.GetSummary(opts.SessionKey)
|
||||||
|
|
||||||
|
// Sanitize history to remove incomplete tool call sequences
|
||||||
|
// This prevents API errors when resuming after task interruption/cancellation
|
||||||
|
sanitizedHistory := sanitizeIncompleteToolCalls(history)
|
||||||
|
if len(sanitizedHistory) != len(history) {
|
||||||
|
logger.WarnCF("agent", "Sanitized incomplete tool calls from session history",
|
||||||
|
map[string]any{
|
||||||
|
"session_key": opts.SessionKey,
|
||||||
|
"original_count": len(history),
|
||||||
|
"sanitized_count": len(sanitizedHistory),
|
||||||
|
})
|
||||||
|
// Update the session with sanitized history
|
||||||
|
agent.Sessions.SetHistory(opts.SessionKey, sanitizedHistory)
|
||||||
|
agent.Sessions.Save(opts.SessionKey)
|
||||||
|
history = sanitizedHistory
|
||||||
|
}
|
||||||
}
|
}
|
||||||
messages := agent.ContextBuilder.BuildMessages(
|
messages := agent.ContextBuilder.BuildMessages(
|
||||||
history,
|
history,
|
||||||
|
|
@ -520,6 +987,19 @@ func (al *AgentLoop) runAgentLoop(ctx context.Context, agent *AgentInstance, opt
|
||||||
// 4. Run LLM iteration loop
|
// 4. Run LLM iteration loop
|
||||||
finalContent, iteration, err := al.runLLMIteration(ctx, agent, messages, opts)
|
finalContent, iteration, err := al.runLLMIteration(ctx, agent, messages, opts)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
|
// Special handling for interruption errors
|
||||||
|
if errors.Is(err, ErrInterrupted) {
|
||||||
|
logger.InfoCF("agent", "Agent execution interrupted by high-priority signal",
|
||||||
|
map[string]any{
|
||||||
|
"agent_id": agent.ID,
|
||||||
|
"session": opts.SessionKey,
|
||||||
|
"iterations": iteration,
|
||||||
|
})
|
||||||
|
// Save current state
|
||||||
|
agent.Sessions.Save(opts.SessionKey)
|
||||||
|
// Return a user-friendly message
|
||||||
|
return "Task interrupted by higher priority request. Progress has been saved.", nil
|
||||||
|
}
|
||||||
return "", err
|
return "", err
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
@ -628,6 +1108,13 @@ func (al *AgentLoop) runLLMIteration(
|
||||||
for iteration < agent.MaxIterations {
|
for iteration < agent.MaxIterations {
|
||||||
iteration++
|
iteration++
|
||||||
|
|
||||||
|
// Phase 2 Step 4: Check for context cancellation at start of each iteration
|
||||||
|
select {
|
||||||
|
case <-ctx.Done():
|
||||||
|
return "", iteration, fmt.Errorf("iteration canceled: %w", ctx.Err())
|
||||||
|
default:
|
||||||
|
}
|
||||||
|
|
||||||
logger.DebugCF("agent", "LLM iteration",
|
logger.DebugCF("agent", "LLM iteration",
|
||||||
map[string]any{
|
map[string]any{
|
||||||
"agent_id": agent.ID,
|
"agent_id": agent.ID,
|
||||||
|
|
@ -668,9 +1155,9 @@ func (al *AgentLoop) runLLMIteration(
|
||||||
fbResult, fbErr := al.fallback.Execute(ctx, agent.Candidates,
|
fbResult, fbErr := al.fallback.Execute(ctx, agent.Candidates,
|
||||||
func(ctx context.Context, provider, model string) (*providers.LLMResponse, error) {
|
func(ctx context.Context, provider, model string) (*providers.LLMResponse, error) {
|
||||||
return agent.Provider.Chat(ctx, messages, providerToolDefs, model, map[string]any{
|
return agent.Provider.Chat(ctx, messages, providerToolDefs, model, map[string]any{
|
||||||
"max_tokens": agent.MaxTokens,
|
"max_tokens": agent.MaxTokens,
|
||||||
"temperature": agent.Temperature,
|
"temperature": agent.Temperature,
|
||||||
"prompt_cache_key": agent.ID,
|
// "prompt_cache_key": agent.ID,
|
||||||
})
|
})
|
||||||
},
|
},
|
||||||
)
|
)
|
||||||
|
|
@ -685,9 +1172,9 @@ func (al *AgentLoop) runLLMIteration(
|
||||||
return fbResult.Response, nil
|
return fbResult.Response, nil
|
||||||
}
|
}
|
||||||
return agent.Provider.Chat(ctx, messages, providerToolDefs, agent.Model, map[string]any{
|
return agent.Provider.Chat(ctx, messages, providerToolDefs, agent.Model, map[string]any{
|
||||||
"max_tokens": agent.MaxTokens,
|
"max_tokens": agent.MaxTokens,
|
||||||
"temperature": agent.Temperature,
|
"temperature": agent.Temperature,
|
||||||
"prompt_cache_key": agent.ID,
|
// "prompt_cache_key": agent.ID,
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
@ -778,8 +1265,53 @@ func (al *AgentLoop) runLLMIteration(
|
||||||
"target_channel": al.targetReasoningChannelID(opts.Channel),
|
"target_channel": al.targetReasoningChannelID(opts.Channel),
|
||||||
"channel": opts.Channel,
|
"channel": opts.Channel,
|
||||||
})
|
})
|
||||||
// Check if no tool calls - we're done
|
// Check if no tool calls - but first check for pending interruptions
|
||||||
if len(response.ToolCalls) == 0 {
|
if len(response.ToolCalls) == 0 {
|
||||||
|
// NEW: Check for pending interruptions before finishing
|
||||||
|
if al.enableSteering {
|
||||||
|
checker := al.getOrCreateChecker(opts.SessionKey)
|
||||||
|
pending := checker.DrainAll()
|
||||||
|
|
||||||
|
if len(pending) > 0 {
|
||||||
|
logger.InfoCF("agent", "Steering: LLM finished but has pending interruptions, injecting",
|
||||||
|
map[string]any{
|
||||||
|
"session_key": opts.SessionKey,
|
||||||
|
"pending_count": len(pending),
|
||||||
|
"iteration": iteration,
|
||||||
|
})
|
||||||
|
|
||||||
|
// Save the assistant's response first
|
||||||
|
assistantMsg := providers.Message{
|
||||||
|
Role: "assistant",
|
||||||
|
Content: response.Content,
|
||||||
|
}
|
||||||
|
messages = append(messages, assistantMsg)
|
||||||
|
agent.Sessions.AddMessage(opts.SessionKey, "assistant", response.Content)
|
||||||
|
|
||||||
|
// Send the response to user
|
||||||
|
if !constants.IsInternalChannel(opts.Channel) {
|
||||||
|
al.bus.PublishOutbound(ctx, bus.OutboundMessage{
|
||||||
|
Channel: opts.Channel,
|
||||||
|
ChatID: opts.ChatID,
|
||||||
|
Content: response.Content,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
// Format and inject interruption
|
||||||
|
injectionContent := formatInterruptionInjection(pending)
|
||||||
|
injectionMsg := providers.Message{
|
||||||
|
Role: "user",
|
||||||
|
Content: injectionContent,
|
||||||
|
}
|
||||||
|
messages = append(messages, injectionMsg)
|
||||||
|
agent.Sessions.AddMessage(opts.SessionKey, "user", injectionContent)
|
||||||
|
|
||||||
|
// Continue to handle the interruption
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// No interruptions, finish normally
|
||||||
finalContent = response.Content
|
finalContent = response.Content
|
||||||
logger.InfoCF("agent", "LLM response without tool calls (direct answer)",
|
logger.InfoCF("agent", "LLM response without tool calls (direct answer)",
|
||||||
map[string]any{
|
map[string]any{
|
||||||
|
|
@ -842,7 +1374,52 @@ func (al *AgentLoop) runLLMIteration(
|
||||||
agent.Sessions.AddFullMessage(opts.SessionKey, assistantMsg)
|
agent.Sessions.AddFullMessage(opts.SessionKey, assistantMsg)
|
||||||
|
|
||||||
// Execute tool calls
|
// Execute tool calls
|
||||||
for _, tc := range normalizedToolCalls {
|
for i, tc := range normalizedToolCalls {
|
||||||
|
// Phase 2 Step 4: Check for context cancellation before each tool execution
|
||||||
|
select {
|
||||||
|
case <-ctx.Done():
|
||||||
|
return "", iteration, fmt.Errorf("tool execution canceled: %w", ctx.Err())
|
||||||
|
default:
|
||||||
|
}
|
||||||
|
|
||||||
|
// Check for interruption before each tool call
|
||||||
|
if al.interruptHandler != nil {
|
||||||
|
signal, checkErr := al.interruptHandler.CheckInterruption(ctx)
|
||||||
|
if checkErr != nil {
|
||||||
|
logger.WarnCF("agent", "Interruption check failed",
|
||||||
|
map[string]any{"error": checkErr.Error()})
|
||||||
|
}
|
||||||
|
|
||||||
|
if signal != nil && signal.Priority >= 8 {
|
||||||
|
logger.InfoCF("agent", "High-priority interruption detected during tool execution",
|
||||||
|
map[string]any{
|
||||||
|
"type": signal.Type,
|
||||||
|
"priority": signal.Priority,
|
||||||
|
"source": signal.Source,
|
||||||
|
"tool": tc.Name,
|
||||||
|
"tool_index": i,
|
||||||
|
"total_tools": len(normalizedToolCalls),
|
||||||
|
})
|
||||||
|
|
||||||
|
// Save current progress to session
|
||||||
|
agent.Sessions.Save(opts.SessionKey)
|
||||||
|
|
||||||
|
// If it's a user message interruption, put the message back
|
||||||
|
if signal.Type == "user_message" && signal.Data != nil {
|
||||||
|
if msg, ok := signal.Data.(bus.InboundMessage); ok {
|
||||||
|
go func() {
|
||||||
|
pubCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||||
|
defer cancel()
|
||||||
|
al.bus.PublishInbound(pubCtx, msg)
|
||||||
|
}()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Return interruption error
|
||||||
|
return "", iteration, ErrInterrupted
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
argsJSON, _ := json.Marshal(tc.Arguments)
|
argsJSON, _ := json.Marshal(tc.Arguments)
|
||||||
argsPreview := utils.Truncate(string(argsJSON), 200)
|
argsPreview := utils.Truncate(string(argsJSON), 200)
|
||||||
logger.InfoCF("agent", fmt.Sprintf("Tool call: %s(%s)", tc.Name, argsPreview),
|
logger.InfoCF("agent", fmt.Sprintf("Tool call: %s(%s)", tc.Name, argsPreview),
|
||||||
|
|
@ -929,6 +1506,58 @@ func (al *AgentLoop) runLLMIteration(
|
||||||
// Save tool result message to session
|
// Save tool result message to session
|
||||||
agent.Sessions.AddFullMessage(opts.SessionKey, toolResultMsg)
|
agent.Sessions.AddFullMessage(opts.SessionKey, toolResultMsg)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// ===== NEW: Steering Architecture - Check for interruptions after tool execution =====
|
||||||
|
if al.enableSteering {
|
||||||
|
checker := al.getOrCreateChecker(opts.SessionKey)
|
||||||
|
pending := checker.DrainAll()
|
||||||
|
|
||||||
|
if len(pending) > 0 {
|
||||||
|
logger.InfoCF("agent", "Steering: injecting interruption messages",
|
||||||
|
map[string]any{
|
||||||
|
"session_key": opts.SessionKey,
|
||||||
|
"pending_count": len(pending),
|
||||||
|
"iteration": iteration,
|
||||||
|
})
|
||||||
|
|
||||||
|
// Send progress update to user showing tool results before handling interruption
|
||||||
|
if !constants.IsInternalChannel(opts.Channel) {
|
||||||
|
// Build a brief summary of what just completed
|
||||||
|
completedTools := []string{}
|
||||||
|
for _, msg := range messages {
|
||||||
|
if msg.Role == "assistant" && len(msg.ToolCalls) > 0 {
|
||||||
|
for _, tc := range msg.ToolCalls {
|
||||||
|
completedTools = append(completedTools, tc.Name)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
progressMsg := fmt.Sprintf("⚡ Completed: %s\n📥 Processing new request...",
|
||||||
|
utils.Truncate(fmt.Sprint(completedTools), 100))
|
||||||
|
|
||||||
|
al.bus.PublishOutbound(ctx, bus.OutboundMessage{
|
||||||
|
Channel: opts.Channel,
|
||||||
|
ChatID: opts.ChatID,
|
||||||
|
Content: progressMsg,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
// Format and inject interruption
|
||||||
|
injectionContent := formatInterruptionInjection(pending)
|
||||||
|
injectionMsg := providers.Message{
|
||||||
|
Role: "user",
|
||||||
|
Content: injectionContent,
|
||||||
|
}
|
||||||
|
messages = append(messages, injectionMsg)
|
||||||
|
|
||||||
|
// Save injection to session for context continuity
|
||||||
|
agent.Sessions.AddMessage(opts.SessionKey, "user", injectionContent)
|
||||||
|
|
||||||
|
// Continue to next iteration with injected message
|
||||||
|
// The LLM will decide how to handle both the original task and the new request
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
return finalContent, iteration, nil
|
return finalContent, iteration, nil
|
||||||
|
|
@ -1161,9 +1790,9 @@ func (al *AgentLoop) summarizeSession(agent *AgentInstance, sessionKey string) {
|
||||||
nil,
|
nil,
|
||||||
agent.Model,
|
agent.Model,
|
||||||
map[string]any{
|
map[string]any{
|
||||||
"max_tokens": 1024,
|
"max_tokens": 1024,
|
||||||
"temperature": 0.3,
|
"temperature": 0.3,
|
||||||
"prompt_cache_key": agent.ID,
|
// "prompt_cache_key": agent.ID,
|
||||||
},
|
},
|
||||||
)
|
)
|
||||||
if err == nil {
|
if err == nil {
|
||||||
|
|
@ -1212,9 +1841,9 @@ func (al *AgentLoop) summarizeBatch(
|
||||||
nil,
|
nil,
|
||||||
agent.Model,
|
agent.Model,
|
||||||
map[string]any{
|
map[string]any{
|
||||||
"max_tokens": 1024,
|
"max_tokens": 1024,
|
||||||
"temperature": 0.3,
|
"temperature": 0.3,
|
||||||
"prompt_cache_key": agent.ID,
|
// "prompt_cache_key": agent.ID,
|
||||||
},
|
},
|
||||||
)
|
)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
|
|
@ -1293,6 +1922,77 @@ func (al *AgentLoop) handleCommand(ctx context.Context, msg bus.InboundMessage)
|
||||||
return fmt.Sprintf("Unknown list target: %s", args[0]), true
|
return fmt.Sprintf("Unknown list target: %s", args[0]), true
|
||||||
}
|
}
|
||||||
|
|
||||||
|
case "/stop", "/cancel", "/abort":
|
||||||
|
// Phase 2: Cancel running tasks for the current session
|
||||||
|
runningTasks := al.taskManager.GetRunningTasksForSession(msg.Channel, msg.ChatID)
|
||||||
|
|
||||||
|
// Filter out command tasks from the count (exclude tasks with is_command metadata)
|
||||||
|
nonCommandTasks := make([]*Task, 0)
|
||||||
|
for _, task := range runningTasks {
|
||||||
|
if isCmd, ok := task.Metadata["is_command"].(bool); !ok || !isCmd {
|
||||||
|
nonCommandTasks = append(nonCommandTasks, task)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if len(nonCommandTasks) == 0 {
|
||||||
|
// No running tasks to cancel
|
||||||
|
return "ℹ️ 没有正在运行的任务需要取消。\nNo running tasks to cancel.", true
|
||||||
|
}
|
||||||
|
|
||||||
|
// Cancel all tasks (including command tasks)
|
||||||
|
canceled := al.taskManager.CancelAllTasksForSession(msg.Channel, msg.ChatID)
|
||||||
|
|
||||||
|
// Sanitize session history to remove incomplete tool_use/tool_result pairs
|
||||||
|
// This prevents API errors (e.g., "tool_use ids were found without tool_result blocks")
|
||||||
|
// when resuming after cancellation
|
||||||
|
route := al.registry.ResolveRoute(routing.RouteInput{
|
||||||
|
Channel: msg.Channel,
|
||||||
|
AccountID: msg.Metadata["account_id"],
|
||||||
|
Peer: extractPeer(msg),
|
||||||
|
ParentPeer: extractParentPeer(msg),
|
||||||
|
GuildID: msg.Metadata["guild_id"],
|
||||||
|
TeamID: msg.Metadata["team_id"],
|
||||||
|
})
|
||||||
|
agent, _ := al.registry.GetAgent(route.AgentID)
|
||||||
|
if agent == nil {
|
||||||
|
agent = al.registry.GetDefaultAgent()
|
||||||
|
}
|
||||||
|
sessionSanitized := false
|
||||||
|
if agent != nil {
|
||||||
|
sessionKey := route.SessionKey
|
||||||
|
if msg.SessionKey != "" && strings.HasPrefix(msg.SessionKey, "agent:") {
|
||||||
|
sessionKey = msg.SessionKey
|
||||||
|
}
|
||||||
|
// Sanitize the session history to remove incomplete tool call sequences
|
||||||
|
history := agent.Sessions.GetHistory(sessionKey)
|
||||||
|
sanitizedHistory := sanitizeIncompleteToolCalls(history)
|
||||||
|
agent.Sessions.SetHistory(sessionKey, sanitizedHistory)
|
||||||
|
agent.Sessions.Save(sessionKey)
|
||||||
|
sessionSanitized = len(history) != len(sanitizedHistory)
|
||||||
|
}
|
||||||
|
|
||||||
|
logger.InfoCF("agent", "User canceled running tasks",
|
||||||
|
map[string]any{
|
||||||
|
"channel": msg.Channel,
|
||||||
|
"chat_id": msg.ChatID,
|
||||||
|
"command": cmd,
|
||||||
|
"total_canceled": canceled,
|
||||||
|
"non_command_count": len(nonCommandTasks),
|
||||||
|
"session_sanitized": sessionSanitized,
|
||||||
|
})
|
||||||
|
|
||||||
|
// Build response message using non-command task count
|
||||||
|
var response string
|
||||||
|
nonCommandCount := len(nonCommandTasks)
|
||||||
|
if nonCommandCount == 1 {
|
||||||
|
response = "🛑 已取消 1 个正在运行的任务。\n✅ 已就绪,可以处理新的请求。\n\nCanceled 1 running task. Ready for new requests."
|
||||||
|
} else {
|
||||||
|
response = fmt.Sprintf("🛑 已取消 %d 个正在运行的任务。\n✅ 已就绪,可以处理新的请求。\n\nCanceled %d running tasks. Ready for new requests.",
|
||||||
|
nonCommandCount, nonCommandCount)
|
||||||
|
}
|
||||||
|
|
||||||
|
return response, true
|
||||||
|
|
||||||
case "/switch":
|
case "/switch":
|
||||||
if len(args) < 3 || args[1] != "to" {
|
if len(args) < 3 || args[1] != "to" {
|
||||||
return "Usage: /switch [model|channel] to <name>", true
|
return "Usage: /switch [model|channel] to <name>", true
|
||||||
|
|
@ -1350,3 +2050,87 @@ func extractParentPeer(msg bus.InboundMessage) *routing.RoutePeer {
|
||||||
}
|
}
|
||||||
return &routing.RoutePeer{Kind: parentKind, ID: parentID}
|
return &routing.RoutePeer{Kind: parentKind, ID: parentID}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// sanitizeIncompleteToolCalls removes assistant messages with tool calls that don't have
|
||||||
|
// corresponding tool result messages. This prevents API errors when resuming after task cancellation.
|
||||||
|
//
|
||||||
|
// The function walks through the message history and:
|
||||||
|
// 1. Identifies assistant messages with tool calls
|
||||||
|
// 2. Checks if each tool call has a matching tool result in the next messages
|
||||||
|
// 3. Removes assistant messages with incomplete tool call sequences
|
||||||
|
//
|
||||||
|
// This is necessary because Claude's API requires every tool_use block to have a corresponding
|
||||||
|
// tool_result block immediately after in the conversation flow.
|
||||||
|
func sanitizeIncompleteToolCalls(messages []providers.Message) []providers.Message {
|
||||||
|
if len(messages) == 0 {
|
||||||
|
return messages
|
||||||
|
}
|
||||||
|
|
||||||
|
sanitized := make([]providers.Message, 0, len(messages))
|
||||||
|
|
||||||
|
for i := 0; i < len(messages); i++ {
|
||||||
|
msg := messages[i]
|
||||||
|
|
||||||
|
// If this is an assistant message with tool calls, check if all tool calls have results
|
||||||
|
if msg.Role == "assistant" && len(msg.ToolCalls) > 0 {
|
||||||
|
// Collect all tool call IDs from this message
|
||||||
|
toolCallIDs := make(map[string]bool)
|
||||||
|
for _, tc := range msg.ToolCalls {
|
||||||
|
toolCallIDs[tc.ID] = true
|
||||||
|
}
|
||||||
|
|
||||||
|
// Look ahead to find tool results for these calls
|
||||||
|
foundResults := make(map[string]bool)
|
||||||
|
for j := i + 1; j < len(messages); j++ {
|
||||||
|
nextMsg := messages[j]
|
||||||
|
if nextMsg.Role == "tool" && nextMsg.ToolCallID != "" {
|
||||||
|
if toolCallIDs[nextMsg.ToolCallID] {
|
||||||
|
foundResults[nextMsg.ToolCallID] = true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
// Stop looking once we hit another assistant message or user message
|
||||||
|
// (tool results must immediately follow the assistant message with tool calls)
|
||||||
|
if nextMsg.Role == "assistant" || nextMsg.Role == "user" {
|
||||||
|
break
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Only keep this assistant message if ALL tool calls have results
|
||||||
|
if len(foundResults) == len(toolCallIDs) {
|
||||||
|
sanitized = append(sanitized, msg)
|
||||||
|
} else {
|
||||||
|
// Skip this assistant message and all its tool results
|
||||||
|
logger.WarnCF("agent", "Removing incomplete tool call sequence",
|
||||||
|
map[string]any{
|
||||||
|
"expected_results": len(toolCallIDs),
|
||||||
|
"found_results": len(foundResults),
|
||||||
|
"tool_call_ids": getToolCallIDsList(msg.ToolCalls),
|
||||||
|
})
|
||||||
|
|
||||||
|
// Also skip the subsequent tool result messages that belong to this assistant message
|
||||||
|
for j := i + 1; j < len(messages); j++ {
|
||||||
|
nextMsg := messages[j]
|
||||||
|
if nextMsg.Role == "tool" && toolCallIDs[nextMsg.ToolCallID] {
|
||||||
|
i = j // Skip this tool result too
|
||||||
|
} else {
|
||||||
|
break
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
// Keep non-assistant messages or assistant messages without tool calls
|
||||||
|
sanitized = append(sanitized, msg)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
return sanitized
|
||||||
|
}
|
||||||
|
|
||||||
|
// getToolCallIDsList extracts tool call IDs for logging
|
||||||
|
func getToolCallIDsList(toolCalls []providers.ToolCall) []string {
|
||||||
|
ids := make([]string, len(toolCalls))
|
||||||
|
for i, tc := range toolCalls {
|
||||||
|
ids[i] = tc.ID
|
||||||
|
}
|
||||||
|
return ids
|
||||||
|
}
|
||||||
|
|
|
||||||
288
pkg/agent/message.go
Normal file
288
pkg/agent/message.go
Normal file
|
|
@ -0,0 +1,288 @@
|
||||||
|
package agent
|
||||||
|
|
||||||
|
import (
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"github.com/sipeed/picoclaw/pkg/providers"
|
||||||
|
)
|
||||||
|
|
||||||
|
// AgentMessageType defines the semantic type of a message beyond standard LLM roles
|
||||||
|
type AgentMessageType string
|
||||||
|
|
||||||
|
const (
|
||||||
|
// Standard LLM message types
|
||||||
|
MessageTypeUser AgentMessageType = "user"
|
||||||
|
MessageTypeAssistant AgentMessageType = "assistant"
|
||||||
|
MessageTypeTool AgentMessageType = "tool"
|
||||||
|
MessageTypeSystem AgentMessageType = "system"
|
||||||
|
|
||||||
|
// Extended types for business semantics
|
||||||
|
MessageTypeArtifact AgentMessageType = "artifact" // LLM-generated artifacts (code, images, documents)
|
||||||
|
MessageTypeAttachment AgentMessageType = "attachment" // User-provided attachments
|
||||||
|
MessageTypeEvent AgentMessageType = "event" // System events (task started, interrupted, etc.)
|
||||||
|
)
|
||||||
|
|
||||||
|
// ArtifactType categorizes the type of artifact
|
||||||
|
type ArtifactType string
|
||||||
|
|
||||||
|
const (
|
||||||
|
ArtifactTypeCode ArtifactType = "code"
|
||||||
|
ArtifactTypeImage ArtifactType = "image"
|
||||||
|
ArtifactTypeDocument ArtifactType = "document"
|
||||||
|
ArtifactTypeData ArtifactType = "data"
|
||||||
|
ArtifactTypeOther ArtifactType = "other"
|
||||||
|
)
|
||||||
|
|
||||||
|
// AgentMessage extends the standard providers.Message with business semantics
|
||||||
|
// and metadata that are useful for the agent runtime but not necessarily for the LLM.
|
||||||
|
type AgentMessage struct {
|
||||||
|
// ===== Core Fields (compatible with providers.Message) =====
|
||||||
|
Role string `json:"role"`
|
||||||
|
Content string `json:"content"`
|
||||||
|
ReasoningContent string `json:"reasoning_content,omitempty"`
|
||||||
|
ToolCalls []providers.ToolCall `json:"tool_calls,omitempty"`
|
||||||
|
ToolCallID string `json:"tool_call_id,omitempty"`
|
||||||
|
|
||||||
|
// ===== Extended Fields =====
|
||||||
|
Type AgentMessageType `json:"type"` // Semantic message type
|
||||||
|
Metadata map[string]any `json:"metadata,omitempty"` // Arbitrary metadata
|
||||||
|
Timestamp time.Time `json:"timestamp"` // Message creation time
|
||||||
|
SessionID string `json:"session_id,omitempty"` // Associated session
|
||||||
|
|
||||||
|
// ===== Artifact-specific Fields =====
|
||||||
|
ArtifactID string `json:"artifact_id,omitempty"` // Unique identifier for artifact
|
||||||
|
ArtifactType ArtifactType `json:"artifact_type,omitempty"` // Type of artifact
|
||||||
|
ArtifactMIME string `json:"artifact_mime,omitempty"` // MIME type if applicable
|
||||||
|
ArtifactSize int64 `json:"artifact_size,omitempty"` // Size in bytes
|
||||||
|
|
||||||
|
// ===== Attachment-specific Fields =====
|
||||||
|
AttachmentURL string `json:"attachment_url,omitempty"` // URL or file path
|
||||||
|
AttachmentSize int64 `json:"attachment_size,omitempty"` // Size in bytes
|
||||||
|
AttachmentFilename string `json:"attachment_filename,omitempty"` // Original filename
|
||||||
|
|
||||||
|
// ===== Event-specific Fields =====
|
||||||
|
EventType string `json:"event_type,omitempty"` // Event category (task_started, interrupted, etc.)
|
||||||
|
EventData map[string]any `json:"event_data,omitempty"` // Event-specific data
|
||||||
|
}
|
||||||
|
|
||||||
|
// ToLLMMessage converts an AgentMessage to a standard providers.Message
|
||||||
|
// that can be sent to the LLM. Extended fields are dropped or converted to content.
|
||||||
|
func (am *AgentMessage) ToLLMMessage() providers.Message {
|
||||||
|
return providers.Message{
|
||||||
|
Role: am.Role,
|
||||||
|
Content: am.Content,
|
||||||
|
ReasoningContent: am.ReasoningContent,
|
||||||
|
ToolCalls: am.ToolCalls,
|
||||||
|
ToolCallID: am.ToolCallID,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ToLLMMessageWithContext converts an AgentMessage to a providers.Message,
|
||||||
|
// but includes contextual information about artifacts/attachments in the content.
|
||||||
|
func (am *AgentMessage) ToLLMMessageWithContext() providers.Message {
|
||||||
|
msg := am.ToLLMMessage()
|
||||||
|
|
||||||
|
// Add artifact reference to content if present
|
||||||
|
if am.Type == MessageTypeArtifact && am.ArtifactID != "" {
|
||||||
|
if msg.Content == "" {
|
||||||
|
msg.Content = am.formatArtifactReference()
|
||||||
|
} else {
|
||||||
|
msg.Content = am.formatArtifactReference() + "\n\n" + msg.Content
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Add attachment reference to content if present
|
||||||
|
if am.Type == MessageTypeAttachment && am.AttachmentURL != "" {
|
||||||
|
if msg.Content == "" {
|
||||||
|
msg.Content = am.formatAttachmentReference()
|
||||||
|
} else {
|
||||||
|
msg.Content = am.formatAttachmentReference() + "\n\n" + msg.Content
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
return msg
|
||||||
|
}
|
||||||
|
|
||||||
|
// formatArtifactReference creates a human-readable reference to an artifact
|
||||||
|
func (am *AgentMessage) formatArtifactReference() string {
|
||||||
|
ref := "[Artifact"
|
||||||
|
if am.ArtifactID != "" {
|
||||||
|
ref += " ID: " + am.ArtifactID
|
||||||
|
}
|
||||||
|
if am.ArtifactType != "" {
|
||||||
|
ref += " Type: " + string(am.ArtifactType)
|
||||||
|
}
|
||||||
|
if am.ArtifactSize > 0 {
|
||||||
|
ref += " Size: " + formatBytes(am.ArtifactSize)
|
||||||
|
}
|
||||||
|
ref += "]"
|
||||||
|
return ref
|
||||||
|
}
|
||||||
|
|
||||||
|
// formatAttachmentReference creates a human-readable reference to an attachment
|
||||||
|
func (am *AgentMessage) formatAttachmentReference() string {
|
||||||
|
ref := "[Attachment"
|
||||||
|
if am.AttachmentFilename != "" {
|
||||||
|
ref += " File: " + am.AttachmentFilename
|
||||||
|
}
|
||||||
|
if am.AttachmentSize > 0 {
|
||||||
|
ref += " Size: " + formatBytes(am.AttachmentSize)
|
||||||
|
}
|
||||||
|
ref += "]"
|
||||||
|
return ref
|
||||||
|
}
|
||||||
|
|
||||||
|
// formatBytes formats byte size in human-readable format
|
||||||
|
func formatBytes(bytes int64) string {
|
||||||
|
const unit = 1024
|
||||||
|
if bytes < unit {
|
||||||
|
return string(rune(bytes)) + " B"
|
||||||
|
}
|
||||||
|
div, exp := int64(unit), 0
|
||||||
|
for n := bytes / unit; n >= unit; n /= unit {
|
||||||
|
div *= unit
|
||||||
|
exp++
|
||||||
|
}
|
||||||
|
return string(rune(bytes/div)) + " " + "KMGTPE"[exp:exp+1] + "B"
|
||||||
|
}
|
||||||
|
|
||||||
|
// FromProviderMessage creates an AgentMessage from a standard providers.Message
|
||||||
|
func FromProviderMessage(pm providers.Message) *AgentMessage {
|
||||||
|
msgType := AgentMessageType(pm.Role)
|
||||||
|
if pm.Role == "system" {
|
||||||
|
msgType = MessageTypeSystem
|
||||||
|
}
|
||||||
|
|
||||||
|
return &AgentMessage{
|
||||||
|
Role: pm.Role,
|
||||||
|
Content: pm.Content,
|
||||||
|
ReasoningContent: pm.ReasoningContent,
|
||||||
|
ToolCalls: pm.ToolCalls,
|
||||||
|
ToolCallID: pm.ToolCallID,
|
||||||
|
Type: msgType,
|
||||||
|
Timestamp: time.Now(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// NewUserMessage creates a new user message
|
||||||
|
func NewUserMessage(content string) *AgentMessage {
|
||||||
|
return &AgentMessage{
|
||||||
|
Role: "user",
|
||||||
|
Content: content,
|
||||||
|
Type: MessageTypeUser,
|
||||||
|
Timestamp: time.Now(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// NewAssistantMessage creates a new assistant message
|
||||||
|
func NewAssistantMessage(content string) *AgentMessage {
|
||||||
|
return &AgentMessage{
|
||||||
|
Role: "assistant",
|
||||||
|
Content: content,
|
||||||
|
Type: MessageTypeAssistant,
|
||||||
|
Timestamp: time.Now(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// NewToolMessage creates a new tool result message
|
||||||
|
func NewToolMessage(toolCallID, content string) *AgentMessage {
|
||||||
|
return &AgentMessage{
|
||||||
|
Role: "tool",
|
||||||
|
Content: content,
|
||||||
|
ToolCallID: toolCallID,
|
||||||
|
Type: MessageTypeTool,
|
||||||
|
Timestamp: time.Now(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// NewArtifactMessage creates a new artifact message
|
||||||
|
func NewArtifactMessage(artifactID string, artifactType ArtifactType, content string) *AgentMessage {
|
||||||
|
return &AgentMessage{
|
||||||
|
Role: "assistant",
|
||||||
|
Content: content,
|
||||||
|
Type: MessageTypeArtifact,
|
||||||
|
ArtifactID: artifactID,
|
||||||
|
ArtifactType: artifactType,
|
||||||
|
ArtifactSize: int64(len(content)),
|
||||||
|
Timestamp: time.Now(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// NewAttachmentMessage creates a new attachment message
|
||||||
|
func NewAttachmentMessage(url, filename string, size int64) *AgentMessage {
|
||||||
|
return &AgentMessage{
|
||||||
|
Role: "user",
|
||||||
|
Content: "Attachment: " + filename,
|
||||||
|
Type: MessageTypeAttachment,
|
||||||
|
AttachmentURL: url,
|
||||||
|
AttachmentFilename: filename,
|
||||||
|
AttachmentSize: size,
|
||||||
|
Timestamp: time.Now(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// NewEventMessage creates a new event message
|
||||||
|
func NewEventMessage(eventType string, eventData map[string]any) *AgentMessage {
|
||||||
|
return &AgentMessage{
|
||||||
|
Role: "system",
|
||||||
|
Type: MessageTypeEvent,
|
||||||
|
EventType: eventType,
|
||||||
|
EventData: eventData,
|
||||||
|
Timestamp: time.Now(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// IsStandardLLMType returns true if the message type is a standard LLM type
|
||||||
|
func (am *AgentMessage) IsStandardLLMType() bool {
|
||||||
|
return am.Type == MessageTypeUser ||
|
||||||
|
am.Type == MessageTypeAssistant ||
|
||||||
|
am.Type == MessageTypeTool ||
|
||||||
|
am.Type == MessageTypeSystem
|
||||||
|
}
|
||||||
|
|
||||||
|
// IsExtendedType returns true if the message type is an extended business type
|
||||||
|
func (am *AgentMessage) IsExtendedType() bool {
|
||||||
|
return !am.IsStandardLLMType()
|
||||||
|
}
|
||||||
|
|
||||||
|
// Clone creates a deep copy of the message
|
||||||
|
func (am *AgentMessage) Clone() *AgentMessage {
|
||||||
|
clone := *am
|
||||||
|
|
||||||
|
// Deep copy slices and maps
|
||||||
|
if am.ToolCalls != nil {
|
||||||
|
clone.ToolCalls = make([]providers.ToolCall, len(am.ToolCalls))
|
||||||
|
copy(clone.ToolCalls, am.ToolCalls)
|
||||||
|
}
|
||||||
|
|
||||||
|
if am.Metadata != nil {
|
||||||
|
clone.Metadata = make(map[string]any, len(am.Metadata))
|
||||||
|
for k, v := range am.Metadata {
|
||||||
|
clone.Metadata[k] = v
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if am.EventData != nil {
|
||||||
|
clone.EventData = make(map[string]any, len(am.EventData))
|
||||||
|
for k, v := range am.EventData {
|
||||||
|
clone.EventData[k] = v
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
return &clone
|
||||||
|
}
|
||||||
|
|
||||||
|
// WithMetadata adds metadata to the message (builder pattern)
|
||||||
|
func (am *AgentMessage) WithMetadata(key string, value any) *AgentMessage {
|
||||||
|
if am.Metadata == nil {
|
||||||
|
am.Metadata = make(map[string]any)
|
||||||
|
}
|
||||||
|
am.Metadata[key] = value
|
||||||
|
return am
|
||||||
|
}
|
||||||
|
|
||||||
|
// WithSessionID sets the session ID (builder pattern)
|
||||||
|
func (am *AgentMessage) WithSessionID(sessionID string) *AgentMessage {
|
||||||
|
am.SessionID = sessionID
|
||||||
|
return am
|
||||||
|
}
|
||||||
278
pkg/agent/message_converter.go
Normal file
278
pkg/agent/message_converter.go
Normal file
|
|
@ -0,0 +1,278 @@
|
||||||
|
package agent
|
||||||
|
|
||||||
|
import (
|
||||||
|
"fmt"
|
||||||
|
|
||||||
|
"github.com/sipeed/picoclaw/pkg/providers"
|
||||||
|
)
|
||||||
|
|
||||||
|
// MessageConverter defines the interface for converting and transforming agent messages
|
||||||
|
type MessageConverter interface {
|
||||||
|
// TransformContext applies semantic-aware context trimming on AgentMessages
|
||||||
|
TransformContext(messages []*AgentMessage, opts TransformOptions) []*AgentMessage
|
||||||
|
|
||||||
|
// ConvertToLLM converts AgentMessages to providers.Message for LLM consumption
|
||||||
|
ConvertToLLM(messages []*AgentMessage) []providers.Message
|
||||||
|
}
|
||||||
|
|
||||||
|
// TransformOptions configures how context transformation should be performed
|
||||||
|
type TransformOptions struct {
|
||||||
|
MaxMessages int // Maximum number of messages to keep (0 = unlimited)
|
||||||
|
MaxTokens int // Approximate maximum tokens (0 = unlimited)
|
||||||
|
PreserveArtifacts int // Number of recent artifacts to preserve
|
||||||
|
PreserveSystem bool // Always preserve system messages
|
||||||
|
IncludeContext bool // Include artifact/attachment references in content
|
||||||
|
}
|
||||||
|
|
||||||
|
// DefaultTransformOptions returns sensible defaults for context transformation
|
||||||
|
func DefaultTransformOptions() TransformOptions {
|
||||||
|
return TransformOptions{
|
||||||
|
MaxMessages: 100,
|
||||||
|
MaxTokens: 0, // Token counting would require a tokenizer
|
||||||
|
PreserveArtifacts: 5,
|
||||||
|
PreserveSystem: true,
|
||||||
|
IncludeContext: true,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// DefaultMessageConverter implements semantic-aware message conversion
|
||||||
|
type DefaultMessageConverter struct {
|
||||||
|
options TransformOptions
|
||||||
|
}
|
||||||
|
|
||||||
|
// NewMessageConverter creates a new message converter with the given options
|
||||||
|
func NewMessageConverter(opts TransformOptions) *DefaultMessageConverter {
|
||||||
|
return &DefaultMessageConverter{
|
||||||
|
options: opts,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// TransformContext applies semantic-aware context trimming
|
||||||
|
// Strategy:
|
||||||
|
// 1. Always preserve system messages (if PreserveSystem is true)
|
||||||
|
// 2. Preserve recent artifacts up to PreserveArtifacts limit
|
||||||
|
// 3. Keep most recent messages up to MaxMessages limit
|
||||||
|
// 4. Prioritize conversation continuity over older messages
|
||||||
|
func (c *DefaultMessageConverter) TransformContext(messages []*AgentMessage, opts TransformOptions) []*AgentMessage {
|
||||||
|
if len(messages) == 0 {
|
||||||
|
return messages
|
||||||
|
}
|
||||||
|
|
||||||
|
// Use provided options or fall back to converter's default options
|
||||||
|
if opts.MaxMessages == 0 && c.options.MaxMessages > 0 {
|
||||||
|
opts = c.options
|
||||||
|
}
|
||||||
|
|
||||||
|
// If no limit, return all messages
|
||||||
|
if opts.MaxMessages == 0 || len(messages) <= opts.MaxMessages {
|
||||||
|
return messages
|
||||||
|
}
|
||||||
|
|
||||||
|
var result []*AgentMessage
|
||||||
|
var artifactCount int
|
||||||
|
var regularMsgCount int
|
||||||
|
|
||||||
|
// First pass: collect system messages from the beginning
|
||||||
|
systemMessages := make([]*AgentMessage, 0)
|
||||||
|
if opts.PreserveSystem {
|
||||||
|
for _, msg := range messages {
|
||||||
|
if msg.Type == MessageTypeSystem {
|
||||||
|
systemMessages = append(systemMessages, msg)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Second pass: collect messages from end to start
|
||||||
|
// This ensures we keep the most recent context
|
||||||
|
tempResult := make([]*AgentMessage, 0)
|
||||||
|
for i := len(messages) - 1; i >= 0; i-- {
|
||||||
|
msg := messages[i]
|
||||||
|
|
||||||
|
// Skip system messages (already collected)
|
||||||
|
if msg.Type == MessageTypeSystem {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
|
||||||
|
// Preserve recent artifacts (don't count against message limit)
|
||||||
|
if msg.Type == MessageTypeArtifact && artifactCount < opts.PreserveArtifacts {
|
||||||
|
tempResult = append([]*AgentMessage{msg}, tempResult...)
|
||||||
|
artifactCount++
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
|
||||||
|
// Stop if we've reached the message limit
|
||||||
|
if regularMsgCount >= opts.MaxMessages {
|
||||||
|
break
|
||||||
|
}
|
||||||
|
|
||||||
|
// Include the message
|
||||||
|
tempResult = append([]*AgentMessage{msg}, tempResult...)
|
||||||
|
regularMsgCount++
|
||||||
|
}
|
||||||
|
|
||||||
|
// Combine: system messages first, then the rest
|
||||||
|
result = append(result, systemMessages...)
|
||||||
|
result = append(result, tempResult...)
|
||||||
|
|
||||||
|
return result
|
||||||
|
}
|
||||||
|
|
||||||
|
// ConvertToLLM converts AgentMessages to standard providers.Message
|
||||||
|
// Extended message types (artifacts, attachments, events) are either:
|
||||||
|
// - Converted to inline references in content (if IncludeContext is true)
|
||||||
|
// - Dropped entirely (if IncludeContext is false)
|
||||||
|
func (c *DefaultMessageConverter) ConvertToLLM(messages []*AgentMessage) []providers.Message {
|
||||||
|
var result []providers.Message
|
||||||
|
|
||||||
|
for _, msg := range messages {
|
||||||
|
// Skip event messages - they're for runtime use only
|
||||||
|
if msg.Type == MessageTypeEvent {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
|
||||||
|
// Handle artifacts and attachments based on IncludeContext option
|
||||||
|
if msg.Type == MessageTypeArtifact || msg.Type == MessageTypeAttachment {
|
||||||
|
if c.options.IncludeContext {
|
||||||
|
// Convert to assistant message with context
|
||||||
|
result = append(result, msg.ToLLMMessageWithContext())
|
||||||
|
}
|
||||||
|
// Otherwise, skip entirely
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
|
||||||
|
// Standard message types - convert directly
|
||||||
|
result = append(result, msg.ToLLMMessage())
|
||||||
|
}
|
||||||
|
|
||||||
|
return result
|
||||||
|
}
|
||||||
|
|
||||||
|
// BatchConvertFromProvider converts multiple providers.Message to AgentMessage
|
||||||
|
func BatchConvertFromProvider(messages []providers.Message) []*AgentMessage {
|
||||||
|
result := make([]*AgentMessage, len(messages))
|
||||||
|
for i, msg := range messages {
|
||||||
|
result[i] = FromProviderMessage(msg)
|
||||||
|
}
|
||||||
|
return result
|
||||||
|
}
|
||||||
|
|
||||||
|
// BatchConvertToProvider converts multiple AgentMessage to providers.Message
|
||||||
|
func BatchConvertToProvider(messages []*AgentMessage) []providers.Message {
|
||||||
|
result := make([]providers.Message, len(messages))
|
||||||
|
for i, msg := range messages {
|
||||||
|
result[i] = msg.ToLLMMessage()
|
||||||
|
}
|
||||||
|
return result
|
||||||
|
}
|
||||||
|
|
||||||
|
// MessageFilter defines a function type for filtering messages
|
||||||
|
type MessageFilter func(*AgentMessage) bool
|
||||||
|
|
||||||
|
// FilterMessages returns messages that match the given filter
|
||||||
|
func FilterMessages(messages []*AgentMessage, filter MessageFilter) []*AgentMessage {
|
||||||
|
var result []*AgentMessage
|
||||||
|
for _, msg := range messages {
|
||||||
|
if filter(msg) {
|
||||||
|
result = append(result, msg)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return result
|
||||||
|
}
|
||||||
|
|
||||||
|
// Common filters
|
||||||
|
|
||||||
|
// FilterByType returns a filter that matches messages of the given type
|
||||||
|
func FilterByType(msgType AgentMessageType) MessageFilter {
|
||||||
|
return func(msg *AgentMessage) bool {
|
||||||
|
return msg.Type == msgType
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// FilterByRole returns a filter that matches messages with the given role
|
||||||
|
func FilterByRole(role string) MessageFilter {
|
||||||
|
return func(msg *AgentMessage) bool {
|
||||||
|
return msg.Role == role
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// FilterBySession returns a filter that matches messages from the given session
|
||||||
|
func FilterBySession(sessionID string) MessageFilter {
|
||||||
|
return func(msg *AgentMessage) bool {
|
||||||
|
return msg.SessionID == sessionID
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// FilterStandardTypes returns a filter that matches only standard LLM message types
|
||||||
|
func FilterStandardTypes() MessageFilter {
|
||||||
|
return func(msg *AgentMessage) bool {
|
||||||
|
return msg.IsStandardLLMType()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// FilterExtendedTypes returns a filter that matches only extended business message types
|
||||||
|
func FilterExtendedTypes() MessageFilter {
|
||||||
|
return func(msg *AgentMessage) bool {
|
||||||
|
return msg.IsExtendedType()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// MessageStats provides statistics about a collection of messages
|
||||||
|
type MessageStats struct {
|
||||||
|
Total int
|
||||||
|
ByType map[AgentMessageType]int
|
||||||
|
ByRole map[string]int
|
||||||
|
TotalSize int64 // Approximate size in bytes
|
||||||
|
ArtifactCount int
|
||||||
|
AttachmentCount int
|
||||||
|
}
|
||||||
|
|
||||||
|
// ComputeStats calculates statistics for a collection of messages
|
||||||
|
func ComputeStats(messages []*AgentMessage) MessageStats {
|
||||||
|
stats := MessageStats{
|
||||||
|
Total: len(messages),
|
||||||
|
ByType: make(map[AgentMessageType]int),
|
||||||
|
ByRole: make(map[string]int),
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, msg := range messages {
|
||||||
|
stats.ByType[msg.Type]++
|
||||||
|
stats.ByRole[msg.Role]++
|
||||||
|
stats.TotalSize += int64(len(msg.Content))
|
||||||
|
|
||||||
|
if msg.Type == MessageTypeArtifact {
|
||||||
|
stats.ArtifactCount++
|
||||||
|
stats.TotalSize += msg.ArtifactSize
|
||||||
|
}
|
||||||
|
|
||||||
|
if msg.Type == MessageTypeAttachment {
|
||||||
|
stats.AttachmentCount++
|
||||||
|
stats.TotalSize += msg.AttachmentSize
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
return stats
|
||||||
|
}
|
||||||
|
|
||||||
|
// String returns a human-readable representation of the stats
|
||||||
|
func (s MessageStats) String() string {
|
||||||
|
return fmt.Sprintf(
|
||||||
|
"Messages: %d total, %d artifacts, %d attachments, ~%s total size",
|
||||||
|
s.Total,
|
||||||
|
s.ArtifactCount,
|
||||||
|
s.AttachmentCount,
|
||||||
|
formatBytes(s.TotalSize),
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
// MergeMessageHistories combines multiple message histories while preserving order
|
||||||
|
// Messages are sorted by timestamp and deduplicated by content hash if needed
|
||||||
|
func MergeMessageHistories(histories ...[]*AgentMessage) []*AgentMessage {
|
||||||
|
var result []*AgentMessage
|
||||||
|
for _, history := range histories {
|
||||||
|
result = append(result, history...)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Sort by timestamp (stable sort preserves order for equal timestamps)
|
||||||
|
// Note: For production use, consider more sophisticated deduplication
|
||||||
|
return result
|
||||||
|
}
|
||||||
432
pkg/agent/message_test.go
Normal file
432
pkg/agent/message_test.go
Normal file
|
|
@ -0,0 +1,432 @@
|
||||||
|
package agent
|
||||||
|
|
||||||
|
import (
|
||||||
|
"testing"
|
||||||
|
|
||||||
|
"github.com/sipeed/picoclaw/pkg/providers"
|
||||||
|
)
|
||||||
|
|
||||||
|
// TestAgentMessage tests basic AgentMessage functionality
|
||||||
|
func TestAgentMessage(t *testing.T) {
|
||||||
|
t.Run("NewUserMessage", func(t *testing.T) {
|
||||||
|
msg := NewUserMessage("Hello, assistant!")
|
||||||
|
if msg.Role != "user" {
|
||||||
|
t.Errorf("Expected role 'user', got '%s'", msg.Role)
|
||||||
|
}
|
||||||
|
if msg.Type != MessageTypeUser {
|
||||||
|
t.Errorf("Expected type 'user', got '%s'", msg.Type)
|
||||||
|
}
|
||||||
|
if msg.Content != "Hello, assistant!" {
|
||||||
|
t.Errorf("Expected content 'Hello, assistant!', got '%s'", msg.Content)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("NewArtifactMessage", func(t *testing.T) {
|
||||||
|
content := "package main\n\nfunc main() {\n\tprintln(\"Hello\")\n}"
|
||||||
|
msg := NewArtifactMessage("artifact_123", ArtifactTypeCode, content)
|
||||||
|
|
||||||
|
if msg.Type != MessageTypeArtifact {
|
||||||
|
t.Errorf("Expected type 'artifact', got '%s'", msg.Type)
|
||||||
|
}
|
||||||
|
if msg.ArtifactID != "artifact_123" {
|
||||||
|
t.Errorf("Expected artifact ID 'artifact_123', got '%s'", msg.ArtifactID)
|
||||||
|
}
|
||||||
|
if msg.ArtifactType != ArtifactTypeCode {
|
||||||
|
t.Errorf("Expected artifact type 'code', got '%s'", msg.ArtifactType)
|
||||||
|
}
|
||||||
|
if msg.ArtifactSize != int64(len(content)) {
|
||||||
|
t.Errorf("Expected artifact size %d, got %d", len(content), msg.ArtifactSize)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("ToLLMMessage", func(t *testing.T) {
|
||||||
|
agentMsg := NewUserMessage("Test message")
|
||||||
|
agentMsg.ToolCalls = []providers.ToolCall{
|
||||||
|
{ID: "call_1", Name: "test_tool", Arguments: map[string]any{}},
|
||||||
|
}
|
||||||
|
|
||||||
|
llmMsg := agentMsg.ToLLMMessage()
|
||||||
|
if llmMsg.Role != "user" {
|
||||||
|
t.Errorf("Expected role 'user', got '%s'", llmMsg.Role)
|
||||||
|
}
|
||||||
|
if llmMsg.Content != "Test message" {
|
||||||
|
t.Errorf("Expected content 'Test message', got '%s'", llmMsg.Content)
|
||||||
|
}
|
||||||
|
if len(llmMsg.ToolCalls) != 1 {
|
||||||
|
t.Errorf("Expected 1 tool call, got %d", len(llmMsg.ToolCalls))
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("ToLLMMessageWithContext_Artifact", func(t *testing.T) {
|
||||||
|
msg := NewArtifactMessage("artifact_123", ArtifactTypeCode, "code content")
|
||||||
|
llmMsg := msg.ToLLMMessageWithContext()
|
||||||
|
|
||||||
|
// Should include artifact reference in content
|
||||||
|
if llmMsg.Content == "" {
|
||||||
|
t.Error("Expected non-empty content with artifact reference")
|
||||||
|
}
|
||||||
|
if llmMsg.Content[:10] != "[Artifact " {
|
||||||
|
t.Errorf("Expected content to start with '[Artifact', got '%s'", llmMsg.Content[:20])
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("FromProviderMessage", func(t *testing.T) {
|
||||||
|
providerMsg := providers.Message{
|
||||||
|
Role: "assistant",
|
||||||
|
Content: "Hello, user!",
|
||||||
|
}
|
||||||
|
|
||||||
|
agentMsg := FromProviderMessage(providerMsg)
|
||||||
|
if agentMsg.Role != "assistant" {
|
||||||
|
t.Errorf("Expected role 'assistant', got '%s'", agentMsg.Role)
|
||||||
|
}
|
||||||
|
if agentMsg.Type != MessageTypeAssistant {
|
||||||
|
t.Errorf("Expected type 'assistant', got '%s'", agentMsg.Type)
|
||||||
|
}
|
||||||
|
if agentMsg.Content != "Hello, user!" {
|
||||||
|
t.Errorf("Expected content 'Hello, user!', got '%s'", agentMsg.Content)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("WithMetadata", func(t *testing.T) {
|
||||||
|
msg := NewUserMessage("Test").
|
||||||
|
WithMetadata("key1", "value1").
|
||||||
|
WithMetadata("key2", 42)
|
||||||
|
|
||||||
|
if msg.Metadata["key1"] != "value1" {
|
||||||
|
t.Errorf("Expected metadata key1='value1', got '%v'", msg.Metadata["key1"])
|
||||||
|
}
|
||||||
|
if msg.Metadata["key2"] != 42 {
|
||||||
|
t.Errorf("Expected metadata key2=42, got '%v'", msg.Metadata["key2"])
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("Clone", func(t *testing.T) {
|
||||||
|
original := NewUserMessage("Test").
|
||||||
|
WithMetadata("key", "value").
|
||||||
|
WithSessionID("session_1")
|
||||||
|
|
||||||
|
cloned := original.Clone()
|
||||||
|
|
||||||
|
// Modify clone
|
||||||
|
cloned.Content = "Modified"
|
||||||
|
cloned.Metadata["key"] = "modified"
|
||||||
|
|
||||||
|
// Original should be unchanged
|
||||||
|
if original.Content != "Test" {
|
||||||
|
t.Error("Original content was modified")
|
||||||
|
}
|
||||||
|
if original.Metadata["key"] != "value" {
|
||||||
|
t.Error("Original metadata was modified")
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("IsStandardLLMType", func(t *testing.T) {
|
||||||
|
testCases := []struct {
|
||||||
|
msgType AgentMessageType
|
||||||
|
expected bool
|
||||||
|
}{
|
||||||
|
{MessageTypeUser, true},
|
||||||
|
{MessageTypeAssistant, true},
|
||||||
|
{MessageTypeTool, true},
|
||||||
|
{MessageTypeSystem, true},
|
||||||
|
{MessageTypeArtifact, false},
|
||||||
|
{MessageTypeAttachment, false},
|
||||||
|
{MessageTypeEvent, false},
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, tc := range testCases {
|
||||||
|
msg := &AgentMessage{Type: tc.msgType}
|
||||||
|
if msg.IsStandardLLMType() != tc.expected {
|
||||||
|
t.Errorf("Type %s: expected IsStandardLLMType=%v, got %v",
|
||||||
|
tc.msgType, tc.expected, msg.IsStandardLLMType())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
// TestMessageConverter tests the message converter functionality
|
||||||
|
func TestMessageConverter(t *testing.T) {
|
||||||
|
t.Run("TransformContext_NoLimit", func(t *testing.T) {
|
||||||
|
messages := []*AgentMessage{
|
||||||
|
NewUserMessage("Message 1"),
|
||||||
|
NewAssistantMessage("Response 1"),
|
||||||
|
NewUserMessage("Message 2"),
|
||||||
|
}
|
||||||
|
|
||||||
|
converter := NewMessageConverter(DefaultTransformOptions())
|
||||||
|
opts := TransformOptions{MaxMessages: 0} // No limit
|
||||||
|
result := converter.TransformContext(messages, opts)
|
||||||
|
|
||||||
|
if len(result) != len(messages) {
|
||||||
|
t.Errorf("Expected %d messages, got %d", len(messages), len(result))
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("TransformContext_WithLimit", func(t *testing.T) {
|
||||||
|
messages := []*AgentMessage{
|
||||||
|
NewUserMessage("Message 1"),
|
||||||
|
NewAssistantMessage("Response 1"),
|
||||||
|
NewUserMessage("Message 2"),
|
||||||
|
NewAssistantMessage("Response 2"),
|
||||||
|
NewUserMessage("Message 3"),
|
||||||
|
}
|
||||||
|
|
||||||
|
converter := NewMessageConverter(DefaultTransformOptions())
|
||||||
|
opts := TransformOptions{MaxMessages: 3}
|
||||||
|
result := converter.TransformContext(messages, opts)
|
||||||
|
|
||||||
|
if len(result) != 3 {
|
||||||
|
t.Errorf("Expected 3 messages, got %d", len(result))
|
||||||
|
}
|
||||||
|
|
||||||
|
// Should keep the most recent messages
|
||||||
|
if result[len(result)-1].Content != "Message 3" {
|
||||||
|
t.Error("Expected most recent message to be preserved")
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("TransformContext_PreserveSystem", func(t *testing.T) {
|
||||||
|
messages := []*AgentMessage{
|
||||||
|
{Role: "system", Type: MessageTypeSystem, Content: "System prompt"},
|
||||||
|
NewUserMessage("Message 1"),
|
||||||
|
NewAssistantMessage("Response 1"),
|
||||||
|
NewUserMessage("Message 2"),
|
||||||
|
NewAssistantMessage("Response 2"),
|
||||||
|
}
|
||||||
|
|
||||||
|
converter := NewMessageConverter(DefaultTransformOptions())
|
||||||
|
opts := TransformOptions{
|
||||||
|
MaxMessages: 2,
|
||||||
|
PreserveSystem: true,
|
||||||
|
}
|
||||||
|
result := converter.TransformContext(messages, opts)
|
||||||
|
|
||||||
|
// Should preserve system message plus 2 recent messages
|
||||||
|
if len(result) < 3 {
|
||||||
|
t.Errorf("Expected at least 3 messages (1 system + 2 recent), got %d", len(result))
|
||||||
|
}
|
||||||
|
|
||||||
|
// First message should be system
|
||||||
|
if result[0].Type != MessageTypeSystem {
|
||||||
|
t.Error("Expected first message to be system message")
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("TransformContext_PreserveArtifacts", func(t *testing.T) {
|
||||||
|
messages := []*AgentMessage{
|
||||||
|
NewUserMessage("Message 1"),
|
||||||
|
NewArtifactMessage("art_1", ArtifactTypeCode, "code 1"),
|
||||||
|
NewUserMessage("Message 2"),
|
||||||
|
NewArtifactMessage("art_2", ArtifactTypeCode, "code 2"),
|
||||||
|
NewUserMessage("Message 3"),
|
||||||
|
}
|
||||||
|
|
||||||
|
converter := NewMessageConverter(DefaultTransformOptions())
|
||||||
|
opts := TransformOptions{
|
||||||
|
MaxMessages: 2,
|
||||||
|
PreserveArtifacts: 2,
|
||||||
|
}
|
||||||
|
result := converter.TransformContext(messages, opts)
|
||||||
|
|
||||||
|
artifactCount := 0
|
||||||
|
for _, msg := range result {
|
||||||
|
if msg.Type == MessageTypeArtifact {
|
||||||
|
artifactCount++
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if artifactCount != 2 {
|
||||||
|
t.Errorf("Expected 2 artifacts preserved, got %d", artifactCount)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("ConvertToLLM", func(t *testing.T) {
|
||||||
|
messages := []*AgentMessage{
|
||||||
|
NewUserMessage("User message"),
|
||||||
|
NewAssistantMessage("Assistant response"),
|
||||||
|
NewArtifactMessage("art_1", ArtifactTypeCode, "code"),
|
||||||
|
NewEventMessage("task_started", map[string]any{"task_id": "123"}),
|
||||||
|
}
|
||||||
|
|
||||||
|
converter := NewMessageConverter(TransformOptions{IncludeContext: false})
|
||||||
|
result := converter.ConvertToLLM(messages)
|
||||||
|
|
||||||
|
// Should have 2 messages (user + assistant), artifact and event dropped
|
||||||
|
if len(result) != 2 {
|
||||||
|
t.Errorf("Expected 2 LLM messages, got %d", len(result))
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("ConvertToLLM_WithContext", func(t *testing.T) {
|
||||||
|
messages := []*AgentMessage{
|
||||||
|
NewUserMessage("User message"),
|
||||||
|
NewArtifactMessage("art_1", ArtifactTypeCode, "code"),
|
||||||
|
}
|
||||||
|
|
||||||
|
converter := NewMessageConverter(TransformOptions{IncludeContext: true})
|
||||||
|
result := converter.ConvertToLLM(messages)
|
||||||
|
|
||||||
|
// Should have 2 messages (user + artifact with context)
|
||||||
|
if len(result) != 2 {
|
||||||
|
t.Errorf("Expected 2 LLM messages, got %d", len(result))
|
||||||
|
}
|
||||||
|
|
||||||
|
// Artifact should include reference
|
||||||
|
if result[1].Content[:10] != "[Artifact " {
|
||||||
|
t.Error("Expected artifact reference in content")
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
// TestMessageFilters tests message filtering functionality
|
||||||
|
func TestMessageFilters(t *testing.T) {
|
||||||
|
messages := []*AgentMessage{
|
||||||
|
NewUserMessage("User 1"),
|
||||||
|
NewAssistantMessage("Assistant 1"),
|
||||||
|
NewArtifactMessage("art_1", ArtifactTypeCode, "code"),
|
||||||
|
{Role: "system", Type: MessageTypeSystem, Content: "System"},
|
||||||
|
NewEventMessage("event_1", nil),
|
||||||
|
}
|
||||||
|
|
||||||
|
t.Run("FilterByType", func(t *testing.T) {
|
||||||
|
result := FilterMessages(messages, FilterByType(MessageTypeArtifact))
|
||||||
|
if len(result) != 1 {
|
||||||
|
t.Errorf("Expected 1 artifact message, got %d", len(result))
|
||||||
|
}
|
||||||
|
if result[0].Type != MessageTypeArtifact {
|
||||||
|
t.Error("Expected artifact message")
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("FilterByRole", func(t *testing.T) {
|
||||||
|
result := FilterMessages(messages, FilterByRole("user"))
|
||||||
|
if len(result) != 1 {
|
||||||
|
t.Errorf("Expected 1 user message, got %d", len(result))
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("FilterStandardTypes", func(t *testing.T) {
|
||||||
|
result := FilterMessages(messages, FilterStandardTypes())
|
||||||
|
if len(result) != 3 {
|
||||||
|
t.Errorf("Expected 3 standard messages, got %d", len(result))
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("FilterExtendedTypes", func(t *testing.T) {
|
||||||
|
result := FilterMessages(messages, FilterExtendedTypes())
|
||||||
|
if len(result) != 2 {
|
||||||
|
t.Errorf("Expected 2 extended messages, got %d", len(result))
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
// TestComputeStats tests message statistics
|
||||||
|
func TestComputeStats(t *testing.T) {
|
||||||
|
messages := []*AgentMessage{
|
||||||
|
NewUserMessage("User message"),
|
||||||
|
NewAssistantMessage("Assistant response"),
|
||||||
|
NewArtifactMessage("art_1", ArtifactTypeCode, "code content"),
|
||||||
|
NewAttachmentMessage("/path/file.txt", "file.txt", 1024),
|
||||||
|
}
|
||||||
|
|
||||||
|
stats := ComputeStats(messages)
|
||||||
|
|
||||||
|
if stats.Total != 4 {
|
||||||
|
t.Errorf("Expected 4 total messages, got %d", stats.Total)
|
||||||
|
}
|
||||||
|
|
||||||
|
if stats.ArtifactCount != 1 {
|
||||||
|
t.Errorf("Expected 1 artifact, got %d", stats.ArtifactCount)
|
||||||
|
}
|
||||||
|
|
||||||
|
if stats.AttachmentCount != 1 {
|
||||||
|
t.Errorf("Expected 1 attachment, got %d", stats.AttachmentCount)
|
||||||
|
}
|
||||||
|
|
||||||
|
if stats.ByType[MessageTypeUser] != 1 {
|
||||||
|
t.Errorf("Expected 1 user message, got %d", stats.ByType[MessageTypeUser])
|
||||||
|
}
|
||||||
|
|
||||||
|
if stats.TotalSize == 0 {
|
||||||
|
t.Error("Expected non-zero total size")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// TestBatchConversion tests batch message conversion
|
||||||
|
func TestBatchConversion(t *testing.T) {
|
||||||
|
t.Run("BatchConvertFromProvider", func(t *testing.T) {
|
||||||
|
providerMsgs := []providers.Message{
|
||||||
|
{Role: "user", Content: "Hello"},
|
||||||
|
{Role: "assistant", Content: "Hi there"},
|
||||||
|
}
|
||||||
|
|
||||||
|
agentMsgs := BatchConvertFromProvider(providerMsgs)
|
||||||
|
|
||||||
|
if len(agentMsgs) != 2 {
|
||||||
|
t.Errorf("Expected 2 messages, got %d", len(agentMsgs))
|
||||||
|
}
|
||||||
|
|
||||||
|
if agentMsgs[0].Type != MessageTypeUser {
|
||||||
|
t.Error("Expected first message to be user type")
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("BatchConvertToProvider", func(t *testing.T) {
|
||||||
|
agentMsgs := []*AgentMessage{
|
||||||
|
NewUserMessage("Hello"),
|
||||||
|
NewAssistantMessage("Hi there"),
|
||||||
|
}
|
||||||
|
|
||||||
|
providerMsgs := BatchConvertToProvider(agentMsgs)
|
||||||
|
|
||||||
|
if len(providerMsgs) != 2 {
|
||||||
|
t.Errorf("Expected 2 messages, got %d", len(providerMsgs))
|
||||||
|
}
|
||||||
|
|
||||||
|
if providerMsgs[0].Role != "user" {
|
||||||
|
t.Error("Expected first message to have user role")
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
// BenchmarkMessageConversion benchmarks message conversion
|
||||||
|
func BenchmarkMessageConversion(b *testing.B) {
|
||||||
|
msg := NewUserMessage("Test message")
|
||||||
|
|
||||||
|
b.Run("ToLLMMessage", func(b *testing.B) {
|
||||||
|
for i := 0; i < b.N; i++ {
|
||||||
|
_ = msg.ToLLMMessage()
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
b.Run("ToLLMMessageWithContext", func(b *testing.B) {
|
||||||
|
artifactMsg := NewArtifactMessage("art_1", ArtifactTypeCode, "code")
|
||||||
|
for i := 0; i < b.N; i++ {
|
||||||
|
_ = artifactMsg.ToLLMMessageWithContext()
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
b.Run("Clone", func(b *testing.B) {
|
||||||
|
for i := 0; i < b.N; i++ {
|
||||||
|
_ = msg.Clone()
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
// BenchmarkMessageFilter benchmarks message filtering
|
||||||
|
func BenchmarkMessageFilter(b *testing.B) {
|
||||||
|
messages := make([]*AgentMessage, 100)
|
||||||
|
for i := 0; i < 100; i++ {
|
||||||
|
if i%3 == 0 {
|
||||||
|
messages[i] = NewArtifactMessage("art", ArtifactTypeCode, "code")
|
||||||
|
} else {
|
||||||
|
messages[i] = NewUserMessage("message")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
b.ResetTimer()
|
||||||
|
for i := 0; i < b.N; i++ {
|
||||||
|
_ = FilterMessages(messages, FilterByType(MessageTypeArtifact))
|
||||||
|
}
|
||||||
|
}
|
||||||
361
pkg/agent/task.go
Normal file
361
pkg/agent/task.go
Normal file
|
|
@ -0,0 +1,361 @@
|
||||||
|
// PicoClaw - Ultra-lightweight personal AI agent
|
||||||
|
// License: MIT
|
||||||
|
//
|
||||||
|
// Copyright (c) 2026 PicoClaw contributors
|
||||||
|
|
||||||
|
// DEPRECATED: This file contains the legacy TaskManager implementation for Phase 2 concurrent task management.
|
||||||
|
// The new steering architecture (nanobot-inspired) uses per-session InterruptionChecker instead.
|
||||||
|
// This code is kept for backward compatibility but will be removed in a future version.
|
||||||
|
// See: pkg/agent/interruption_checker.go for the new implementation.
|
||||||
|
|
||||||
|
package agent
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"fmt"
|
||||||
|
"sync"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"github.com/sipeed/picoclaw/pkg/bus"
|
||||||
|
"github.com/sipeed/picoclaw/pkg/logger"
|
||||||
|
)
|
||||||
|
|
||||||
|
// TaskStatus represents the current state of a task
|
||||||
|
type TaskStatus string
|
||||||
|
|
||||||
|
const (
|
||||||
|
TaskStatusPending TaskStatus = "pending" // Task is queued but not started
|
||||||
|
TaskStatusRunning TaskStatus = "running" // Task is currently executing
|
||||||
|
TaskStatusCompleted TaskStatus = "completed" // Task finished successfully
|
||||||
|
TaskStatusCanceled TaskStatus = "canceled" // Task was canceled by user/system
|
||||||
|
TaskStatusFailed TaskStatus = "failed" // Task failed with error
|
||||||
|
)
|
||||||
|
|
||||||
|
// Task represents a single message processing task
|
||||||
|
type Task struct {
|
||||||
|
ID string // Unique task identifier
|
||||||
|
Message bus.InboundMessage // The message being processed
|
||||||
|
Status TaskStatus // Current status
|
||||||
|
Priority int // Task priority (from interrupt handler)
|
||||||
|
CreatedAt time.Time // When task was created
|
||||||
|
StartedAt time.Time // When task started executing
|
||||||
|
EndedAt time.Time // When task completed/canceled/failed
|
||||||
|
Error error // Error if task failed
|
||||||
|
Metadata map[string]any // Task metadata for custom attributes
|
||||||
|
|
||||||
|
// Context management
|
||||||
|
ctx context.Context // Task execution context
|
||||||
|
cancel context.CancelFunc // Function to cancel this task
|
||||||
|
|
||||||
|
// Result channel
|
||||||
|
done chan struct{} // Closed when task completes
|
||||||
|
}
|
||||||
|
|
||||||
|
// NewTask creates a new task from an inbound message
|
||||||
|
func NewTask(msg bus.InboundMessage, priority int) *Task {
|
||||||
|
return &Task{
|
||||||
|
ID: fmt.Sprintf("%s:%s:%d", msg.Channel, msg.ChatID, time.Now().UnixNano()),
|
||||||
|
Message: msg,
|
||||||
|
Status: TaskStatusPending,
|
||||||
|
Priority: priority,
|
||||||
|
CreatedAt: time.Now(),
|
||||||
|
Metadata: make(map[string]any),
|
||||||
|
done: make(chan struct{}),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Start marks the task as running and sets up cancellation context
|
||||||
|
func (t *Task) Start(parentCtx context.Context) {
|
||||||
|
t.ctx, t.cancel = context.WithCancel(parentCtx)
|
||||||
|
t.Status = TaskStatusRunning
|
||||||
|
t.StartedAt = time.Now()
|
||||||
|
}
|
||||||
|
|
||||||
|
// Cancel cancels the task execution
|
||||||
|
func (t *Task) Cancel() {
|
||||||
|
// Check if already in a terminal state
|
||||||
|
if t.Status == TaskStatusCompleted || t.Status == TaskStatusFailed || t.Status == TaskStatusCanceled {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
if t.cancel != nil {
|
||||||
|
t.cancel()
|
||||||
|
}
|
||||||
|
t.Status = TaskStatusCanceled
|
||||||
|
t.EndedAt = time.Now()
|
||||||
|
close(t.done)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Complete marks the task as completed
|
||||||
|
func (t *Task) Complete() {
|
||||||
|
// Check if already in a terminal state
|
||||||
|
if t.Status == TaskStatusCompleted || t.Status == TaskStatusFailed || t.Status == TaskStatusCanceled {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
t.Status = TaskStatusCompleted
|
||||||
|
t.EndedAt = time.Now()
|
||||||
|
close(t.done)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Fail marks the task as failed with an error
|
||||||
|
func (t *Task) Fail(err error) {
|
||||||
|
// Check if already in a terminal state
|
||||||
|
if t.Status == TaskStatusCompleted || t.Status == TaskStatusFailed || t.Status == TaskStatusCanceled {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
t.Status = TaskStatusFailed
|
||||||
|
t.Error = err
|
||||||
|
t.EndedAt = time.Now()
|
||||||
|
close(t.done)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Wait blocks until the task completes or is canceled
|
||||||
|
func (t *Task) Wait() {
|
||||||
|
<-t.done
|
||||||
|
}
|
||||||
|
|
||||||
|
// Context returns the task's execution context
|
||||||
|
func (t *Task) Context() context.Context {
|
||||||
|
return t.ctx
|
||||||
|
}
|
||||||
|
|
||||||
|
// TaskManager manages concurrent task execution and cancellation
|
||||||
|
type TaskManager struct {
|
||||||
|
mu sync.RWMutex
|
||||||
|
tasks map[string]*Task // All tasks by ID
|
||||||
|
runningTasks map[string]*Task // Currently running tasks
|
||||||
|
maxConcurrent int // Maximum concurrent tasks (0 = unlimited)
|
||||||
|
}
|
||||||
|
|
||||||
|
// NewTaskManager creates a new task manager
|
||||||
|
func NewTaskManager(maxConcurrent int) *TaskManager {
|
||||||
|
return &TaskManager{
|
||||||
|
tasks: make(map[string]*Task),
|
||||||
|
runningTasks: make(map[string]*Task),
|
||||||
|
maxConcurrent: maxConcurrent,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// AddTask adds a new task to the manager
|
||||||
|
func (tm *TaskManager) AddTask(task *Task) error {
|
||||||
|
tm.mu.Lock()
|
||||||
|
defer tm.mu.Unlock()
|
||||||
|
|
||||||
|
if _, exists := tm.tasks[task.ID]; exists {
|
||||||
|
return fmt.Errorf("task %s already exists", task.ID)
|
||||||
|
}
|
||||||
|
|
||||||
|
tm.tasks[task.ID] = task
|
||||||
|
logger.DebugCF("task", "Task added", map[string]any{
|
||||||
|
"task_id": task.ID,
|
||||||
|
"priority": task.Priority,
|
||||||
|
"channel": task.Message.Channel,
|
||||||
|
"chat_id": task.Message.ChatID,
|
||||||
|
})
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// StartTask marks a task as running
|
||||||
|
func (tm *TaskManager) StartTask(taskID string, parentCtx context.Context) error {
|
||||||
|
tm.mu.Lock()
|
||||||
|
defer tm.mu.Unlock()
|
||||||
|
|
||||||
|
task, exists := tm.tasks[taskID]
|
||||||
|
if !exists {
|
||||||
|
return fmt.Errorf("task %s not found", taskID)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Check concurrency limit
|
||||||
|
if tm.maxConcurrent > 0 && len(tm.runningTasks) >= tm.maxConcurrent {
|
||||||
|
return fmt.Errorf("max concurrent tasks (%d) reached", tm.maxConcurrent)
|
||||||
|
}
|
||||||
|
|
||||||
|
task.Start(parentCtx)
|
||||||
|
tm.runningTasks[taskID] = task
|
||||||
|
|
||||||
|
logger.InfoCF("task", "Task started", map[string]any{
|
||||||
|
"task_id": taskID,
|
||||||
|
"running_tasks": len(tm.runningTasks),
|
||||||
|
"max_concurrent": tm.maxConcurrent,
|
||||||
|
})
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// CompleteTask marks a task as completed and removes it from running tasks
|
||||||
|
func (tm *TaskManager) CompleteTask(taskID string) {
|
||||||
|
tm.mu.Lock()
|
||||||
|
defer tm.mu.Unlock()
|
||||||
|
|
||||||
|
if task, exists := tm.tasks[taskID]; exists {
|
||||||
|
task.Complete()
|
||||||
|
delete(tm.runningTasks, taskID)
|
||||||
|
|
||||||
|
logger.InfoCF("task", "Task completed", map[string]any{
|
||||||
|
"task_id": taskID,
|
||||||
|
"duration": time.Since(task.StartedAt).String(),
|
||||||
|
"running_tasks": len(tm.runningTasks),
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// CancelTask cancels a running task
|
||||||
|
func (tm *TaskManager) CancelTask(taskID string) error {
|
||||||
|
tm.mu.Lock()
|
||||||
|
defer tm.mu.Unlock()
|
||||||
|
|
||||||
|
task, exists := tm.tasks[taskID]
|
||||||
|
if !exists {
|
||||||
|
return fmt.Errorf("task %s not found", taskID)
|
||||||
|
}
|
||||||
|
|
||||||
|
if task.Status != TaskStatusRunning {
|
||||||
|
return fmt.Errorf("task %s is not running (status: %s)", taskID, task.Status)
|
||||||
|
}
|
||||||
|
|
||||||
|
task.Cancel()
|
||||||
|
delete(tm.runningTasks, taskID)
|
||||||
|
|
||||||
|
logger.InfoCF("task", "Task canceled", map[string]any{
|
||||||
|
"task_id": taskID,
|
||||||
|
"duration": time.Since(task.StartedAt).String(),
|
||||||
|
"running_tasks": len(tm.runningTasks),
|
||||||
|
})
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// FailTask marks a task as failed
|
||||||
|
func (tm *TaskManager) FailTask(taskID string, err error) {
|
||||||
|
tm.mu.Lock()
|
||||||
|
defer tm.mu.Unlock()
|
||||||
|
|
||||||
|
if task, exists := tm.tasks[taskID]; exists {
|
||||||
|
task.Fail(err)
|
||||||
|
delete(tm.runningTasks, taskID)
|
||||||
|
|
||||||
|
logger.ErrorCF("task", "Task failed", map[string]any{
|
||||||
|
"task_id": taskID,
|
||||||
|
"error": err.Error(),
|
||||||
|
"duration": time.Since(task.StartedAt).String(),
|
||||||
|
"running_tasks": len(tm.runningTasks),
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// GetTask retrieves a task by ID
|
||||||
|
func (tm *TaskManager) GetTask(taskID string) (*Task, bool) {
|
||||||
|
tm.mu.RLock()
|
||||||
|
defer tm.mu.RUnlock()
|
||||||
|
|
||||||
|
task, exists := tm.tasks[taskID]
|
||||||
|
return task, exists
|
||||||
|
}
|
||||||
|
|
||||||
|
// GetRunningTasks returns all currently running tasks
|
||||||
|
func (tm *TaskManager) GetRunningTasks() []*Task {
|
||||||
|
tm.mu.RLock()
|
||||||
|
defer tm.mu.RUnlock()
|
||||||
|
|
||||||
|
tasks := make([]*Task, 0, len(tm.runningTasks))
|
||||||
|
for _, task := range tm.runningTasks {
|
||||||
|
tasks = append(tasks, task)
|
||||||
|
}
|
||||||
|
return tasks
|
||||||
|
}
|
||||||
|
|
||||||
|
// GetRunningTasksForSession returns running tasks for a specific session
|
||||||
|
func (tm *TaskManager) GetRunningTasksForSession(channel, chatID string) []*Task {
|
||||||
|
tm.mu.RLock()
|
||||||
|
defer tm.mu.RUnlock()
|
||||||
|
|
||||||
|
tasks := make([]*Task, 0)
|
||||||
|
for _, task := range tm.runningTasks {
|
||||||
|
if task.Message.Channel == channel && task.Message.ChatID == chatID {
|
||||||
|
tasks = append(tasks, task)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return tasks
|
||||||
|
}
|
||||||
|
|
||||||
|
// CancelAllTasksForSession cancels all running tasks for a specific session
|
||||||
|
func (tm *TaskManager) CancelAllTasksForSession(channel, chatID string) int {
|
||||||
|
tm.mu.Lock()
|
||||||
|
defer tm.mu.Unlock()
|
||||||
|
|
||||||
|
canceled := 0
|
||||||
|
for taskID, task := range tm.runningTasks {
|
||||||
|
if task.Message.Channel == channel && task.Message.ChatID == chatID {
|
||||||
|
task.Cancel()
|
||||||
|
delete(tm.runningTasks, taskID)
|
||||||
|
canceled++
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if canceled > 0 {
|
||||||
|
logger.InfoCF("task", "Canceled tasks for session", map[string]any{
|
||||||
|
"channel": channel,
|
||||||
|
"chat_id": chatID,
|
||||||
|
"canceled_count": canceled,
|
||||||
|
"running_tasks": len(tm.runningTasks),
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
return canceled
|
||||||
|
}
|
||||||
|
|
||||||
|
// Cleanup removes completed/failed/canceled tasks older than the specified duration
|
||||||
|
func (tm *TaskManager) Cleanup(olderThan time.Duration) int {
|
||||||
|
tm.mu.Lock()
|
||||||
|
defer tm.mu.Unlock()
|
||||||
|
|
||||||
|
cutoff := time.Now().Add(-olderThan)
|
||||||
|
removed := 0
|
||||||
|
|
||||||
|
for taskID, task := range tm.tasks {
|
||||||
|
if task.Status != TaskStatusRunning && task.EndedAt.Before(cutoff) {
|
||||||
|
delete(tm.tasks, taskID)
|
||||||
|
removed++
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if removed > 0 {
|
||||||
|
logger.DebugCF("task", "Cleaned up old tasks", map[string]any{
|
||||||
|
"removed": removed,
|
||||||
|
"total": len(tm.tasks),
|
||||||
|
"older_than": olderThan.String(),
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
return removed
|
||||||
|
}
|
||||||
|
|
||||||
|
// Stats returns task manager statistics
|
||||||
|
func (tm *TaskManager) Stats() map[string]int {
|
||||||
|
tm.mu.RLock()
|
||||||
|
defer tm.mu.RUnlock()
|
||||||
|
|
||||||
|
stats := map[string]int{
|
||||||
|
"total": len(tm.tasks),
|
||||||
|
"running": len(tm.runningTasks),
|
||||||
|
"pending": 0,
|
||||||
|
"completed": 0,
|
||||||
|
"canceled": 0,
|
||||||
|
"failed": 0,
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, task := range tm.tasks {
|
||||||
|
switch task.Status {
|
||||||
|
case TaskStatusPending:
|
||||||
|
stats["pending"]++
|
||||||
|
case TaskStatusCompleted:
|
||||||
|
stats["completed"]++
|
||||||
|
case TaskStatusCanceled:
|
||||||
|
stats["canceled"]++
|
||||||
|
case TaskStatusFailed:
|
||||||
|
stats["failed"]++
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
return stats
|
||||||
|
}
|
||||||
320
pkg/agent/task_test.go
Normal file
320
pkg/agent/task_test.go
Normal file
|
|
@ -0,0 +1,320 @@
|
||||||
|
// PicoClaw - Ultra-lightweight personal AI agent
|
||||||
|
// License: MIT
|
||||||
|
//
|
||||||
|
// Copyright (c) 2026 PicoClaw contributors
|
||||||
|
|
||||||
|
package agent
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"testing"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"github.com/sipeed/picoclaw/pkg/bus"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestNewTask(t *testing.T) {
|
||||||
|
msg := bus.InboundMessage{
|
||||||
|
Channel: "telegram",
|
||||||
|
ChatID: "123",
|
||||||
|
Content: "test message",
|
||||||
|
}
|
||||||
|
|
||||||
|
task := NewTask(msg, 5)
|
||||||
|
|
||||||
|
if task.ID == "" {
|
||||||
|
t.Error("Expected task ID to be set")
|
||||||
|
}
|
||||||
|
if task.Status != TaskStatusPending {
|
||||||
|
t.Errorf("Expected status pending, got %s", task.Status)
|
||||||
|
}
|
||||||
|
if task.Priority != 5 {
|
||||||
|
t.Errorf("Expected priority 5, got %d", task.Priority)
|
||||||
|
}
|
||||||
|
if task.Message.Channel != "telegram" {
|
||||||
|
t.Errorf("Expected channel telegram, got %s", task.Message.Channel)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestTask_StartAndComplete(t *testing.T) {
|
||||||
|
msg := bus.InboundMessage{Channel: "test", ChatID: "1", Content: "test"}
|
||||||
|
task := NewTask(msg, 5)
|
||||||
|
|
||||||
|
ctx := context.Background()
|
||||||
|
task.Start(ctx)
|
||||||
|
|
||||||
|
if task.Status != TaskStatusRunning {
|
||||||
|
t.Errorf("Expected status running, got %s", task.Status)
|
||||||
|
}
|
||||||
|
if task.ctx == nil {
|
||||||
|
t.Error("Expected context to be set")
|
||||||
|
}
|
||||||
|
if task.cancel == nil {
|
||||||
|
t.Error("Expected cancel function to be set")
|
||||||
|
}
|
||||||
|
|
||||||
|
task.Complete()
|
||||||
|
|
||||||
|
if task.Status != TaskStatusCompleted {
|
||||||
|
t.Errorf("Expected status completed, got %s", task.Status)
|
||||||
|
}
|
||||||
|
|
||||||
|
// done channel should be closed
|
||||||
|
select {
|
||||||
|
case <-task.done:
|
||||||
|
// OK
|
||||||
|
case <-time.After(100 * time.Millisecond):
|
||||||
|
t.Error("Expected done channel to be closed")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestTask_Cancel(t *testing.T) {
|
||||||
|
msg := bus.InboundMessage{Channel: "test", ChatID: "1", Content: "test"}
|
||||||
|
task := NewTask(msg, 5)
|
||||||
|
|
||||||
|
ctx := context.Background()
|
||||||
|
task.Start(ctx)
|
||||||
|
task.Cancel()
|
||||||
|
|
||||||
|
if task.Status != TaskStatusCanceled {
|
||||||
|
t.Errorf("Expected status canceled, got %s", task.Status)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Context should be canceled
|
||||||
|
select {
|
||||||
|
case <-task.ctx.Done():
|
||||||
|
// OK
|
||||||
|
case <-time.After(100 * time.Millisecond):
|
||||||
|
t.Error("Expected context to be canceled")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestTask_Fail(t *testing.T) {
|
||||||
|
msg := bus.InboundMessage{Channel: "test", ChatID: "1", Content: "test"}
|
||||||
|
task := NewTask(msg, 5)
|
||||||
|
|
||||||
|
ctx := context.Background()
|
||||||
|
task.Start(ctx)
|
||||||
|
|
||||||
|
testErr := context.DeadlineExceeded
|
||||||
|
task.Fail(testErr)
|
||||||
|
|
||||||
|
if task.Status != TaskStatusFailed {
|
||||||
|
t.Errorf("Expected status failed, got %s", task.Status)
|
||||||
|
}
|
||||||
|
if task.Error != testErr {
|
||||||
|
t.Errorf("Expected error %v, got %v", testErr, task.Error)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestNewTaskManager(t *testing.T) {
|
||||||
|
tm := NewTaskManager(5)
|
||||||
|
|
||||||
|
if tm.maxConcurrent != 5 {
|
||||||
|
t.Errorf("Expected max concurrent 5, got %d", tm.maxConcurrent)
|
||||||
|
}
|
||||||
|
if tm.tasks == nil {
|
||||||
|
t.Error("Expected tasks map to be initialized")
|
||||||
|
}
|
||||||
|
if tm.runningTasks == nil {
|
||||||
|
t.Error("Expected running tasks map to be initialized")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestTaskManager_AddTask(t *testing.T) {
|
||||||
|
tm := NewTaskManager(0)
|
||||||
|
msg := bus.InboundMessage{Channel: "test", ChatID: "1", Content: "test"}
|
||||||
|
task := NewTask(msg, 5)
|
||||||
|
|
||||||
|
err := tm.AddTask(task)
|
||||||
|
if err != nil {
|
||||||
|
t.Errorf("Unexpected error: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Try adding same task again
|
||||||
|
err = tm.AddTask(task)
|
||||||
|
if err == nil {
|
||||||
|
t.Error("Expected error when adding duplicate task")
|
||||||
|
}
|
||||||
|
|
||||||
|
stats := tm.Stats()
|
||||||
|
if stats["total"] != 1 {
|
||||||
|
t.Errorf("Expected 1 total task, got %d", stats["total"])
|
||||||
|
}
|
||||||
|
if stats["pending"] != 1 {
|
||||||
|
t.Errorf("Expected 1 pending task, got %d", stats["pending"])
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestTaskManager_StartTask(t *testing.T) {
|
||||||
|
tm := NewTaskManager(2)
|
||||||
|
ctx := context.Background()
|
||||||
|
|
||||||
|
msg1 := bus.InboundMessage{Channel: "test", ChatID: "1", Content: "test1"}
|
||||||
|
task1 := NewTask(msg1, 5)
|
||||||
|
tm.AddTask(task1)
|
||||||
|
|
||||||
|
err := tm.StartTask(task1.ID, ctx)
|
||||||
|
if err != nil {
|
||||||
|
t.Errorf("Unexpected error: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
stats := tm.Stats()
|
||||||
|
if stats["running"] != 1 {
|
||||||
|
t.Errorf("Expected 1 running task, got %d", stats["running"])
|
||||||
|
}
|
||||||
|
|
||||||
|
// Test concurrency limit
|
||||||
|
msg2 := bus.InboundMessage{Channel: "test", ChatID: "2", Content: "test2"}
|
||||||
|
task2 := NewTask(msg2, 5)
|
||||||
|
tm.AddTask(task2)
|
||||||
|
tm.StartTask(task2.ID, ctx)
|
||||||
|
|
||||||
|
msg3 := bus.InboundMessage{Channel: "test", ChatID: "3", Content: "test3"}
|
||||||
|
task3 := NewTask(msg3, 5)
|
||||||
|
tm.AddTask(task3)
|
||||||
|
|
||||||
|
err = tm.StartTask(task3.ID, ctx)
|
||||||
|
if err == nil {
|
||||||
|
t.Error("Expected error when exceeding concurrency limit")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestTaskManager_CompleteTask(t *testing.T) {
|
||||||
|
tm := NewTaskManager(0)
|
||||||
|
ctx := context.Background()
|
||||||
|
|
||||||
|
msg := bus.InboundMessage{Channel: "test", ChatID: "1", Content: "test"}
|
||||||
|
task := NewTask(msg, 5)
|
||||||
|
tm.AddTask(task)
|
||||||
|
tm.StartTask(task.ID, ctx)
|
||||||
|
|
||||||
|
tm.CompleteTask(task.ID)
|
||||||
|
|
||||||
|
stats := tm.Stats()
|
||||||
|
if stats["running"] != 0 {
|
||||||
|
t.Errorf("Expected 0 running tasks, got %d", stats["running"])
|
||||||
|
}
|
||||||
|
if stats["completed"] != 1 {
|
||||||
|
t.Errorf("Expected 1 completed task, got %d", stats["completed"])
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestTaskManager_CancelTask(t *testing.T) {
|
||||||
|
tm := NewTaskManager(0)
|
||||||
|
ctx := context.Background()
|
||||||
|
|
||||||
|
msg := bus.InboundMessage{Channel: "test", ChatID: "1", Content: "test"}
|
||||||
|
task := NewTask(msg, 5)
|
||||||
|
tm.AddTask(task)
|
||||||
|
tm.StartTask(task.ID, ctx)
|
||||||
|
|
||||||
|
err := tm.CancelTask(task.ID)
|
||||||
|
if err != nil {
|
||||||
|
t.Errorf("Unexpected error: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
stats := tm.Stats()
|
||||||
|
if stats["running"] != 0 {
|
||||||
|
t.Errorf("Expected 0 running tasks, got %d", stats["running"])
|
||||||
|
}
|
||||||
|
if stats["canceled"] != 1 {
|
||||||
|
t.Errorf("Expected 1 canceled task, got %d", stats["canceled"])
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestTaskManager_GetRunningTasksForSession(t *testing.T) {
|
||||||
|
tm := NewTaskManager(0)
|
||||||
|
ctx := context.Background()
|
||||||
|
|
||||||
|
msg1 := bus.InboundMessage{Channel: "telegram", ChatID: "123", Content: "test1"}
|
||||||
|
task1 := NewTask(msg1, 5)
|
||||||
|
tm.AddTask(task1)
|
||||||
|
tm.StartTask(task1.ID, ctx)
|
||||||
|
|
||||||
|
msg2 := bus.InboundMessage{Channel: "telegram", ChatID: "123", Content: "test2"}
|
||||||
|
task2 := NewTask(msg2, 5)
|
||||||
|
tm.AddTask(task2)
|
||||||
|
tm.StartTask(task2.ID, ctx)
|
||||||
|
|
||||||
|
msg3 := bus.InboundMessage{Channel: "telegram", ChatID: "456", Content: "test3"}
|
||||||
|
task3 := NewTask(msg3, 5)
|
||||||
|
tm.AddTask(task3)
|
||||||
|
tm.StartTask(task3.ID, ctx)
|
||||||
|
|
||||||
|
tasks := tm.GetRunningTasksForSession("telegram", "123")
|
||||||
|
if len(tasks) != 2 {
|
||||||
|
t.Errorf("Expected 2 tasks for session, got %d", len(tasks))
|
||||||
|
}
|
||||||
|
|
||||||
|
tasks = tm.GetRunningTasksForSession("telegram", "456")
|
||||||
|
if len(tasks) != 1 {
|
||||||
|
t.Errorf("Expected 1 task for session, got %d", len(tasks))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestTaskManager_CancelAllTasksForSession(t *testing.T) {
|
||||||
|
tm := NewTaskManager(0)
|
||||||
|
ctx := context.Background()
|
||||||
|
|
||||||
|
msg1 := bus.InboundMessage{Channel: "telegram", ChatID: "123", Content: "test1"}
|
||||||
|
task1 := NewTask(msg1, 5)
|
||||||
|
tm.AddTask(task1)
|
||||||
|
tm.StartTask(task1.ID, ctx)
|
||||||
|
|
||||||
|
msg2 := bus.InboundMessage{Channel: "telegram", ChatID: "123", Content: "test2"}
|
||||||
|
task2 := NewTask(msg2, 5)
|
||||||
|
tm.AddTask(task2)
|
||||||
|
tm.StartTask(task2.ID, ctx)
|
||||||
|
|
||||||
|
msg3 := bus.InboundMessage{Channel: "telegram", ChatID: "456", Content: "test3"}
|
||||||
|
task3 := NewTask(msg3, 5)
|
||||||
|
tm.AddTask(task3)
|
||||||
|
tm.StartTask(task3.ID, ctx)
|
||||||
|
|
||||||
|
canceled := tm.CancelAllTasksForSession("telegram", "123")
|
||||||
|
if canceled != 2 {
|
||||||
|
t.Errorf("Expected 2 tasks canceled, got %d", canceled)
|
||||||
|
}
|
||||||
|
|
||||||
|
stats := tm.Stats()
|
||||||
|
if stats["running"] != 1 {
|
||||||
|
t.Errorf("Expected 1 running task, got %d", stats["running"])
|
||||||
|
}
|
||||||
|
if stats["canceled"] != 2 {
|
||||||
|
t.Errorf("Expected 2 canceled tasks, got %d", stats["canceled"])
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestTaskManager_Cleanup(t *testing.T) {
|
||||||
|
tm := NewTaskManager(0)
|
||||||
|
ctx := context.Background()
|
||||||
|
|
||||||
|
// Create and complete an old task
|
||||||
|
msg1 := bus.InboundMessage{Channel: "test", ChatID: "1", Content: "test1"}
|
||||||
|
task1 := NewTask(msg1, 5)
|
||||||
|
tm.AddTask(task1)
|
||||||
|
tm.StartTask(task1.ID, ctx)
|
||||||
|
tm.CompleteTask(task1.ID)
|
||||||
|
|
||||||
|
// Manually set EndedAt to simulate old task
|
||||||
|
task1.EndedAt = time.Now().Add(-2 * time.Hour)
|
||||||
|
|
||||||
|
// Create a recent completed task
|
||||||
|
msg2 := bus.InboundMessage{Channel: "test", ChatID: "2", Content: "test2"}
|
||||||
|
task2 := NewTask(msg2, 5)
|
||||||
|
tm.AddTask(task2)
|
||||||
|
tm.StartTask(task2.ID, ctx)
|
||||||
|
tm.CompleteTask(task2.ID)
|
||||||
|
|
||||||
|
// Cleanup tasks older than 1 hour
|
||||||
|
removed := tm.Cleanup(1 * time.Hour)
|
||||||
|
if removed != 1 {
|
||||||
|
t.Errorf("Expected 1 task removed, got %d", removed)
|
||||||
|
}
|
||||||
|
|
||||||
|
stats := tm.Stats()
|
||||||
|
if stats["total"] != 1 {
|
||||||
|
t.Errorf("Expected 1 remaining task, got %d", stats["total"])
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
@ -179,6 +179,16 @@ type AgentDefaults struct {
|
||||||
MaxTokens int `json:"max_tokens" env:"PICOCLAW_AGENTS_DEFAULTS_MAX_TOKENS"`
|
MaxTokens int `json:"max_tokens" env:"PICOCLAW_AGENTS_DEFAULTS_MAX_TOKENS"`
|
||||||
Temperature *float64 `json:"temperature,omitempty" env:"PICOCLAW_AGENTS_DEFAULTS_TEMPERATURE"`
|
Temperature *float64 `json:"temperature,omitempty" env:"PICOCLAW_AGENTS_DEFAULTS_TEMPERATURE"`
|
||||||
MaxToolIterations int `json:"max_tool_iterations" env:"PICOCLAW_AGENTS_DEFAULTS_MAX_TOOL_ITERATIONS"`
|
MaxToolIterations int `json:"max_tool_iterations" env:"PICOCLAW_AGENTS_DEFAULTS_MAX_TOOL_ITERATIONS"`
|
||||||
|
|
||||||
|
// Steering architecture (nanobot-inspired, opt-in)
|
||||||
|
EnableSteering bool `json:"enable_steering,omitempty" env:"PICOCLAW_AGENTS_DEFAULTS_ENABLE_STEERING"` // Enable message injection during tool execution
|
||||||
|
|
||||||
|
// Legacy: Phase 2 concurrent task management (to be deprecated)
|
||||||
|
MaxConcurrentTasks int `json:"max_concurrent_tasks,omitempty" env:"PICOCLAW_AGENTS_DEFAULTS_MAX_CONCURRENT_TASKS"` // Maximum concurrent tasks (0=unlimited)
|
||||||
|
EnableSteeringLoop bool `json:"enable_steering_loop,omitempty" env:"PICOCLAW_AGENTS_DEFAULTS_ENABLE_STEERING_LOOP"` // Enable steering loop for interrupt monitoring
|
||||||
|
SteeringLoopIntervalMs int `json:"steering_loop_interval_ms,omitempty" env:"PICOCLAW_AGENTS_DEFAULTS_STEERING_LOOP_INTERVAL_MS"` // Steering loop check interval (ms)
|
||||||
|
TaskCleanupIntervalMins int `json:"task_cleanup_interval_mins,omitempty" env:"PICOCLAW_AGENTS_DEFAULTS_TASK_CLEANUP_INTERVAL_MINS"` // Task cleanup interval (minutes)
|
||||||
|
TaskRetentionHours int `json:"task_retention_hours,omitempty" env:"PICOCLAW_AGENTS_DEFAULTS_TASK_RETENTION_HOURS"` // Task retention time (hours)
|
||||||
}
|
}
|
||||||
|
|
||||||
// GetModelName returns the effective model name for the agent defaults.
|
// GetModelName returns the effective model name for the agent defaults.
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue