Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 13 additions & 1 deletion api/pkg/org/application/dispatch/dispatcher.go
Original file line number Diff line number Diff line change
Expand Up @@ -46,12 +46,16 @@ import (
// embedded Queue and focuses on the event-side fan-out.
type Dispatcher struct {
store *store.Store
queue *activation.Queue
queue ActivationQueue
logger *slog.Logger
outbound map[transport.Kind]streaming.Outbound
processorRunner ProcessorRunner
}

type ActivationQueue interface {
Enqueue(orgID string, agentID orgchart.NodeID, trigger activation.Trigger)
}

// ProcessorRunner is the late-bound execution arm that turns an Event
// into the Processor outputs its Topic feeds. application/processing.Runner
// satisfies it; declared here (not imported) so dispatch does not depend
Expand Down Expand Up @@ -95,6 +99,14 @@ func (d *Dispatcher) RegisterProcessorRunner(r ProcessorRunner) {
d.processorRunner = r
}

// RegisterActivationQueue replaces the in-memory activation queue. Production
// uses this to install restart-safe delivery without changing dispatch callers.
func (d *Dispatcher) RegisterActivationQueue(q ActivationQueue) {
if q != nil {
d.queue = q
}
}

// DispatchHire fires a hire-time activation for a freshly-created AI
// Worker. Returns immediately; the activation runs on a goroutine with
// its own background context — independent of the HTTP request that
Expand Down
5 changes: 5 additions & 0 deletions api/pkg/org/application/lifecycle/hire_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -60,6 +60,8 @@ func TestCreate_CreatesBotAndReconciles(t *testing.T) {
t.Parallel()
st := memory.New()
svc := newHireService(st)
cleaner := &recordingAgentDeliveryCleaner{}
svc.AgentDelivery = cleaner
ctx := context.Background()

boss, _ := orgchart.NewNode("w-boss", "# Eng", nil, hireClock(), "org-test")
Expand All @@ -79,6 +81,9 @@ func TestCreate_CreatesBotAndReconciles(t *testing.T) {
if res.Node.AgentID != "app-agent" {
t.Fatalf("agent app id = %q, want app-agent", res.Node.AgentID)
}
if !cleaner.restored || cleaner.orgID != "org-test" || cleaner.agentID != "w-new" {
t.Fatalf("agent delivery was not restored for recreated worker")
}
if _, err := st.Nodes.Get(ctx, "org-test", "w-new"); err != nil {
t.Fatalf("bot not persisted: %v", err)
}
Expand Down
29 changes: 24 additions & 5 deletions api/pkg/org/application/lifecycle/lifecycle.go
Original file line number Diff line number Diff line change
Expand Up @@ -97,14 +97,20 @@ type OrgReconciler interface {
Reconcile(ctx context.Context, orgID string) error
}

type AgentDeliveryLifecycle interface {
CleanupAgent(ctx context.Context, orgID string, agentID orgchart.NodeID) error
RestoreAgent(orgID string, agentID orgchart.NodeID)
}

// Service composes the node-lifecycle operations the REST/MCP layers
// drive. All fields are required; pass nil HelixRuntime only in tests
// that don't need the Helix-side teardown.
type Service struct {
Store *store.Store
Helix HelixRuntime
Agents AgentCreator
Logger *slog.Logger
Store *store.Store
Helix HelixRuntime
Agents AgentCreator
Logger *slog.Logger
AgentDelivery AgentDeliveryLifecycle

// Nodes is the node-mutation service Create delegates the row creation
// to, so the base-tool union and id minting are shared with the
Expand Down Expand Up @@ -305,6 +311,9 @@ func (s *Service) Create(ctx context.Context, orgID string, p CreateParams) (Cre
}
}

if s.AgentDelivery != nil {
s.AgentDelivery.RestoreAgent(orgID, id)
}
if p.DeferActivation {
return CreateResult{Node: node}, nil
}
Expand Down Expand Up @@ -392,7 +401,7 @@ func (s *Service) ReconcileAgentLinks(ctx context.Context, orgID string) error {
// events themselves are intentionally left behind as an audit trail; only
// the Topic row is dropped. A configured Helix project is not node-owned and
// survives deletion; only its reference to the deleted default agent is unset.
func (s *Service) Delete(ctx context.Context, orgID string, id orgchart.NodeID) error {
func (s *Service) Delete(ctx context.Context, orgID string, id orgchart.NodeID) (err error) {
if id == "" {
return errors.New("node id is empty")
}
Expand All @@ -403,6 +412,16 @@ func (s *Service) Delete(ctx context.Context, orgID string, id orgchart.NodeID)
if err != nil {
return fmt.Errorf("get node %q: %w", id, err)
}
if s.AgentDelivery != nil {
if err := s.AgentDelivery.CleanupAgent(ctx, orgID, id); err != nil {
return fmt.Errorf("cleanup agent delivery %q: %w", id, err)
}
defer func() {
if err != nil {
s.AgentDelivery.RestoreAgent(orgID, id)
}
}()
}

// Capture the deleted Node's managers AND reports BEFORE deletion: the
// reporting lines cascade-drop with the node row, so the List* calls
Expand Down
62 changes: 62 additions & 0 deletions api/pkg/org/application/lifecycle/lifecycle_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ package lifecycle_test

import (
"context"
"errors"
"testing"
"time"

Expand All @@ -16,6 +17,24 @@ import (
orggorm "github.com/helixml/helix/api/pkg/org/infrastructure/persistence/gorm"
)

type recordingAgentDeliveryCleaner struct {
orgID string
agentID orgchart.NodeID
restored bool
}

func (c *recordingAgentDeliveryCleaner) CleanupAgent(_ context.Context, orgID string, agentID orgchart.NodeID) error {
c.orgID = orgID
c.agentID = agentID
return nil
}

func (c *recordingAgentDeliveryCleaner) RestoreAgent(orgID string, agentID orgchart.NodeID) {
c.orgID = orgID
c.agentID = agentID
c.restored = true
}

// TestDelete_RemovesBotsTranscript pins the regression behind "we still
// see s-transcript-w-ai-1 and s-transcript-w-test-ai even though those
// bots are gone": the Delete cascade tore down subscriptions, runtime
Expand Down Expand Up @@ -315,6 +334,49 @@ func TestDelete_Guards(t *testing.T) {
}
}

func TestDelete_CleansUpAgentDelivery(t *testing.T) {
t.Parallel()
ctx := context.Background()
st := orggorm.GetOrgTestDB(t)
const orgID = "org-delete-delivery"
seedBot(t, st, orgID, "w-delete")

cleaner := &recordingAgentDeliveryCleaner{}
if err := (&lifecycle.Service{Store: st, AgentDelivery: cleaner}).Delete(ctx, orgID, "w-delete"); err != nil {
t.Fatalf("Delete: %v", err)
}
if cleaner.orgID != orgID || cleaner.agentID != "w-delete" {
t.Fatalf("cleaned delivery for (%q, %q), want (%q, %q)", cleaner.orgID, cleaner.agentID, orgID, "w-delete")
}
}

func TestDelete_RestoresAgentDeliveryWhenRuntimeDeleteFails(t *testing.T) {
t.Parallel()
ctx := context.Background()
st := orggorm.GetOrgTestDB(t)
const orgID = "org-delete-delivery-failure"
bot, err := orgchart.NewNode("w-delete", "# w-delete", nil, time.Now().UTC(), orgID)
if err != nil {
t.Fatalf("new bot: %v", err)
}
bot.AgentID = "app-delete"
if err := st.Nodes.Create(ctx, bot); err != nil {
t.Fatalf("create bot: %v", err)
}

cleaner := &recordingAgentDeliveryCleaner{}
svc := &lifecycle.Service{Store: st, Helix: &lifecycleRuntime{linkedErr: errors.New("delete failed")}, AgentDelivery: cleaner}
if err := svc.Delete(ctx, orgID, bot.ID); err == nil {
t.Fatal("Delete should fail")
}
if !cleaner.restored {
t.Fatal("agent delivery was not restored after failed deletion")
}
if _, err := st.Nodes.Get(ctx, orgID, bot.ID); err != nil {
t.Fatalf("bot should survive failed deletion: %v", err)
}
}

// TestDelete_MissingBotErrorsWithNoSideEffects pins that deleting a bot
// that doesn't exist errors at the get-guard and leaves the graph alone —
// no bystander is swept up by a no-op delete.
Expand Down
Loading