mirror of
https://github.com/hanzoai/agent.git
synced 2026-09-11 21:08:51 +00:00
zeekay
2899aeaa9f
Consolidate all Hanzo agent infrastructure into one repo: - control-plane/: Go control plane server (was agents-new) - sdk/: Go, Python, TypeScript SDKs - deployments/: Docker, Kubernetes, Helm configs - Dockerfile.sdk: Python SDK image (ghcr.io/hanzoai/agent:latest) - CI builds both SDK + control-plane images No more agentfield — everything is hanzoai/agent.
240 lines
7.1 KiB
Go
240 lines
7.1 KiB
Go
package handlers
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strings"
|
|
"sync"
|
|
"testing"
|
|
|
|
"github.com/hanzoai/agents/control-plane/pkg/types"
|
|
|
|
"github.com/gin-gonic/gin"
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
type memoryStorageStub struct {
|
|
mu sync.Mutex
|
|
store map[string]*types.Memory
|
|
events []*types.MemoryChangeEvent
|
|
published []types.MemoryChangeEvent
|
|
setErr error
|
|
eventErr error
|
|
}
|
|
|
|
func newMemoryStorageStub() *memoryStorageStub {
|
|
return &memoryStorageStub{store: make(map[string]*types.Memory)}
|
|
}
|
|
|
|
func (m *memoryStorageStub) composite(scope, scopeID, key string) string {
|
|
return scope + "|" + scopeID + "|" + key
|
|
}
|
|
|
|
func (m *memoryStorageStub) SetMemory(ctx context.Context, memory *types.Memory) error {
|
|
m.mu.Lock()
|
|
defer m.mu.Unlock()
|
|
if m.setErr != nil {
|
|
return m.setErr
|
|
}
|
|
k := m.composite(memory.Scope, memory.ScopeID, memory.Key)
|
|
m.store[k] = memory
|
|
return nil
|
|
}
|
|
|
|
func (m *memoryStorageStub) GetMemory(ctx context.Context, scope, scopeID, key string) (*types.Memory, error) {
|
|
m.mu.Lock()
|
|
defer m.mu.Unlock()
|
|
k := m.composite(scope, scopeID, key)
|
|
mem, ok := m.store[k]
|
|
if !ok {
|
|
return nil, errors.New("not found")
|
|
}
|
|
return mem, nil
|
|
}
|
|
|
|
func (m *memoryStorageStub) DeleteMemory(ctx context.Context, scope, scopeID, key string) error {
|
|
m.mu.Lock()
|
|
defer m.mu.Unlock()
|
|
k := m.composite(scope, scopeID, key)
|
|
if _, ok := m.store[k]; !ok {
|
|
return errors.New("not found")
|
|
}
|
|
delete(m.store, k)
|
|
return nil
|
|
}
|
|
|
|
func (m *memoryStorageStub) ListMemory(ctx context.Context, scope, scopeID string) ([]*types.Memory, error) {
|
|
m.mu.Lock()
|
|
defer m.mu.Unlock()
|
|
var result []*types.Memory
|
|
prefix := scope + "|" + scopeID + "|"
|
|
for k, mem := range m.store {
|
|
if strings.HasPrefix(k, prefix) {
|
|
result = append(result, mem)
|
|
}
|
|
}
|
|
return result, nil
|
|
}
|
|
|
|
func (m *memoryStorageStub) StoreEvent(ctx context.Context, event *types.MemoryChangeEvent) error {
|
|
m.mu.Lock()
|
|
defer m.mu.Unlock()
|
|
if m.eventErr != nil {
|
|
return m.eventErr
|
|
}
|
|
m.events = append(m.events, event)
|
|
return nil
|
|
}
|
|
|
|
func (m *memoryStorageStub) PublishMemoryChange(ctx context.Context, event types.MemoryChangeEvent) error {
|
|
m.mu.Lock()
|
|
defer m.mu.Unlock()
|
|
m.published = append(m.published, event)
|
|
return nil
|
|
}
|
|
|
|
func (m *memoryStorageStub) SetVector(ctx context.Context, record *types.VectorRecord) error {
|
|
return nil
|
|
}
|
|
|
|
func (m *memoryStorageStub) GetVector(ctx context.Context, scope, scopeID, key string) (*types.VectorRecord, error) {
|
|
return nil, nil
|
|
}
|
|
|
|
func (m *memoryStorageStub) DeleteVector(ctx context.Context, scope, scopeID, key string) error {
|
|
return nil
|
|
}
|
|
|
|
func (m *memoryStorageStub) DeleteVectorsByPrefix(ctx context.Context, scope, scopeID, prefix string) (int, error) {
|
|
return 0, nil
|
|
}
|
|
|
|
func (m *memoryStorageStub) SimilaritySearch(ctx context.Context, scope, scopeID string, queryEmbedding []float32, topK int, filters map[string]interface{}) ([]*types.VectorSearchResult, error) {
|
|
return []*types.VectorSearchResult{}, nil
|
|
}
|
|
|
|
func TestSetMemoryHandler_StoresMemoryAndEvent(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
|
|
storage := newMemoryStorageStub()
|
|
router := gin.New()
|
|
router.POST("/memory/set", SetMemoryHandler(storage))
|
|
|
|
body := `{"key":"alpha","data":{"value":1}}`
|
|
req := httptest.NewRequest(http.MethodPost, "/memory/set", strings.NewReader(body))
|
|
req.Header.Set("Content-Type", "application/json")
|
|
req.Header.Set("X-Workflow-ID", "wf-123")
|
|
req.Header.Set("X-Agent-Node-ID", "agent-1")
|
|
req.Header.Set("X-Actor-ID", "user-42")
|
|
|
|
resp := httptest.NewRecorder()
|
|
router.ServeHTTP(resp, req)
|
|
|
|
require.Equal(t, http.StatusOK, resp.Code)
|
|
|
|
var memory types.Memory
|
|
require.NoError(t, json.Unmarshal(resp.Body.Bytes(), &memory))
|
|
require.Equal(t, "alpha", memory.Key)
|
|
require.Equal(t, "workflow", memory.Scope)
|
|
require.Equal(t, "wf-123", memory.ScopeID)
|
|
|
|
stored, err := storage.GetMemory(context.Background(), "workflow", "wf-123", "alpha")
|
|
require.NoError(t, err)
|
|
require.Equal(t, "alpha", stored.Key)
|
|
|
|
require.Len(t, storage.events, 1)
|
|
require.Equal(t, "set", storage.events[0].Action)
|
|
require.Equal(t, "agent-1", storage.events[0].Metadata.AgentID)
|
|
require.Equal(t, "user-42", storage.events[0].Metadata.ActorID)
|
|
require.Equal(t, "wf-123", storage.events[0].Metadata.WorkflowID)
|
|
require.Len(t, storage.published, 1)
|
|
require.Equal(t, "set", storage.published[0].Action)
|
|
}
|
|
|
|
func TestGetMemoryHandler_HierarchicalLookup(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
|
|
storage := newMemoryStorageStub()
|
|
memData, _ := json.Marshal(map[string]any{"v": 99})
|
|
_ = storage.SetMemory(context.Background(), &types.Memory{
|
|
Scope: "session",
|
|
ScopeID: "session-1",
|
|
Key: "beta",
|
|
Data: memData,
|
|
})
|
|
|
|
router := gin.New()
|
|
router.POST("/memory/get", GetMemoryHandler(storage))
|
|
|
|
body := `{"key":"beta"}`
|
|
req := httptest.NewRequest(http.MethodPost, "/memory/get", strings.NewReader(body))
|
|
req.Header.Set("Content-Type", "application/json")
|
|
req.Header.Set("X-Session-ID", "session-1")
|
|
|
|
resp := httptest.NewRecorder()
|
|
router.ServeHTTP(resp, req)
|
|
|
|
require.Equal(t, http.StatusOK, resp.Code)
|
|
|
|
var memory types.Memory
|
|
require.NoError(t, json.Unmarshal(resp.Body.Bytes(), &memory))
|
|
require.Equal(t, "session", memory.Scope)
|
|
require.Equal(t, "session-1", memory.ScopeID)
|
|
}
|
|
|
|
func TestDeleteMemoryHandler_RemovesEntry(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
|
|
storage := newMemoryStorageStub()
|
|
_ = storage.SetMemory(context.Background(), &types.Memory{Scope: "global", ScopeID: "global", Key: "gamma", Data: json.RawMessage(`"value"`)})
|
|
|
|
router := gin.New()
|
|
router.POST("/memory/delete", DeleteMemoryHandler(storage))
|
|
|
|
body := `{"key":"gamma","scope":"global"}`
|
|
req := httptest.NewRequest(http.MethodPost, "/memory/delete", strings.NewReader(body))
|
|
req.Header.Set("Content-Type", "application/json")
|
|
req.Header.Set("X-Agent-Node-ID", "agent-2")
|
|
req.Header.Set("X-Actor-ID", "user-99")
|
|
|
|
resp := httptest.NewRecorder()
|
|
router.ServeHTTP(resp, req)
|
|
|
|
require.Equal(t, http.StatusNoContent, resp.Code)
|
|
|
|
_, err := storage.GetMemory(context.Background(), "global", "global", "gamma")
|
|
require.Error(t, err)
|
|
require.Len(t, storage.events, 1)
|
|
require.Equal(t, "delete", storage.events[0].Action)
|
|
require.Equal(t, "agent-2", storage.events[0].Metadata.AgentID)
|
|
require.Equal(t, "user-99", storage.events[0].Metadata.ActorID)
|
|
require.Len(t, storage.published, 1)
|
|
require.Equal(t, "delete", storage.published[0].Action)
|
|
}
|
|
|
|
func TestListMemoryHandler_ReturnsMemories(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
|
|
storage := newMemoryStorageStub()
|
|
_ = storage.SetMemory(context.Background(), &types.Memory{Scope: "workflow", ScopeID: "wf-1", Key: "k1", Data: json.RawMessage(`1`)})
|
|
_ = storage.SetMemory(context.Background(), &types.Memory{Scope: "workflow", ScopeID: "wf-1", Key: "k2", Data: json.RawMessage(`2`)})
|
|
|
|
router := gin.New()
|
|
router.GET("/memory/list", ListMemoryHandler(storage))
|
|
|
|
req := httptest.NewRequest(http.MethodGet, "/memory/list?scope=workflow", nil)
|
|
req.Header.Set("X-Workflow-ID", "wf-1")
|
|
|
|
resp := httptest.NewRecorder()
|
|
router.ServeHTTP(resp, req)
|
|
|
|
require.Equal(t, http.StatusOK, resp.Code)
|
|
|
|
var memories []types.Memory
|
|
require.NoError(t, json.Unmarshal(resp.Body.Bytes(), &memories))
|
|
require.Len(t, memories, 2)
|
|
}
|