diff --git a/apps/daemon/internal/cli/agent_discovery.go b/apps/daemon/internal/cli/agent_discovery.go index 9c28fb45d..83c5d2bae 100644 --- a/apps/daemon/internal/cli/agent_discovery.go +++ b/apps/daemon/internal/cli/agent_discovery.go @@ -1,8 +1,11 @@ package cli import ( + "bytes" "context" "fmt" + "sync" + "github.com/MiniMax-AI/OpenAgentCore/apps/daemon/internal/agent" "github.com/MiniMax-AI/OpenAgentCore/apps/daemon/internal/paths" ) @@ -25,18 +28,35 @@ func discoverAgentCLIs(parent context.Context, rc *runContext, profile string, d if err != nil { return nil, err } - var out agentCLIDiscovery - available := false - for _, declaration := range declarations { + // Preflights are independent and sit on every cold start, so they run + // concurrently; their output is replayed in declaration order. + type probe struct { + runtime *agent.Runtime + stdout, stderr bytes.Buffer + } + probes := make([]probe, len(declarations)) + var running sync.WaitGroup + for i, declaration := range declarations { if rc.installedKinds != nil && !rc.installedKinds[declaration.Kind] { continue } - runtime := declaration.Discover(parent, agent.DiscoveryOptions{Profile: profile, StateRoot: stateRoot, RuntimeRoot: runtimeRoot, Stdout: rc.stdout, Stderr: rc.stderr}, declaration.Info()) - if runtime == nil { + running.Go(func() { + p := &probes[i] + p.runtime = declaration.Discover(parent, agent.DiscoveryOptions{Profile: profile, StateRoot: stateRoot, RuntimeRoot: runtimeRoot, Stdout: &p.stdout, Stderr: &p.stderr}, declaration.Info()) + }) + } + running.Wait() + var out agentCLIDiscovery + available := false + for i := range probes { + p := &probes[i] + _, _ = rc.stdout.Write(p.stdout.Bytes()) + _, _ = rc.stderr.Write(p.stderr.Bytes()) + if p.runtime == nil { continue } - out = append(out, discoveredHarness{declaration, *runtime}) - available = available || runtime.Info.Available + out = append(out, discoveredHarness{declarations[i], *p.runtime}) + available = available || p.runtime.Info.Available } if err := parent.Err(); err != nil { return out, err diff --git a/apps/daemon/internal/cli/native_discovery_test.go b/apps/daemon/internal/cli/native_discovery_test.go index e3894f84d..389ecdbe3 100644 --- a/apps/daemon/internal/cli/native_discovery_test.go +++ b/apps/daemon/internal/cli/native_discovery_test.go @@ -3,9 +3,13 @@ package cli import ( "context" "errors" + "fmt" "io" "reflect" + "strings" + "sync" "testing" + "time" "github.com/MiniMax-AI/OpenAgentCore/apps/daemon/internal/agent" "github.com/MiniMax-AI/OpenAgentCore/internal/agentdaemon/proto" @@ -17,31 +21,44 @@ func TestDiscoveryAndRegistration(t *testing.T) { runtimeRoot, stateRoot := t.TempDir(), t.TempDir() t.Setenv("OAC_RUNTIME_HOME", runtimeRoot) t.Setenv("OAC_RUNTIME_STATE_DIRECTORY", stateRoot) - called := []string{} + rc := &runContext{stdout: io.Discard, stderr: io.Discard} + expected := []string{"claude_sdk", "codex", "mcode"} + if selected != "" { + rc.installedKinds = map[string]bool{selected: true} + expected = []string{selected} + } + // Each probe waits for all selected probes to start, so a sequential + // discovery fails here instead of passing slowly. + var started sync.WaitGroup + started.Add(len(expected)) + all := make(chan struct{}) + go func() { started.Wait(); close(all) }() declarations := append([]agent.Declaration(nil), harnessDeclarations...) for i := range declarations { declarations[i].Discover = func(_ context.Context, options agent.DiscoveryOptions, info proto.SupportedAgentKind) *agent.Runtime { if options.RuntimeRoot != runtimeRoot || options.StateRoot != stateRoot { - t.Fatal("discovery lost explicit directory context") + t.Error("discovery lost explicit directory context") + } + started.Done() + select { + case <-all: + case <-time.After(5 * time.Second): + t.Error("preflights did not run concurrently") } - called = append(called, info.Kind) + fmt.Fprintln(options.Stdout, info.Kind) info.Available = true info.Version = "test" return &agent.Runtime{Info: info} } } - rc := &runContext{stdout: io.Discard, stderr: io.Discard} - expected := []string{"claude_sdk", "codex", "mcode"} - if selected != "" { - rc.installedKinds = map[string]bool{selected: true} - expected = []string{selected} - } + var stdout strings.Builder + rc.stdout = &stdout discovery, err := discoverAgentCLIs(t.Context(), rc, "default", declarations) if err != nil { t.Fatal(err) } - if !reflect.DeepEqual(called, expected) { - t.Fatalf("probes %v, want %v", called, expected) + if got := strings.Fields(stdout.String()); !reflect.DeepEqual(got, expected) { + t.Fatalf("preflight output %v, want %v", got, expected) } registry := agent.NewRegistry() registerAgentKinds(registry, discovery) diff --git a/services/core/internal/execution/runtime_initialization.go b/services/core/internal/execution/runtime_initialization.go index 3b9762901..2dd4f9d86 100644 --- a/services/core/internal/execution/runtime_initialization.go +++ b/services/core/internal/execution/runtime_initialization.go @@ -33,6 +33,7 @@ func (w *Worker) runEnvironmentInitializations(ctx context.Context) error { return ctx.Err() case id := <-done: delete(active, id) + case <-w.dispatcher.Registry.Announced(): case <-ticker.C: } rows, err := w.dispatcher.SessionsReader.ListEnvironmentInitializations(ctx, cursor) diff --git a/services/core/internal/runtimegateway/capabilities_test.go b/services/core/internal/runtimegateway/capabilities_test.go index a6c26c69b..54baa1774 100644 --- a/services/core/internal/runtimegateway/capabilities_test.go +++ b/services/core/internal/runtimegateway/capabilities_test.go @@ -48,3 +48,25 @@ func TestInvalidHeartbeatDiscardsPreviousDeclarationAndClosesTransport(t *testin t.Fatal("old declaration survived invalid heartbeat") } } + +func TestFirstHeartbeatAnnouncesAgentKinds(t *testing.T) { + registry := NewRegistry() + session := newTestSession(newFakeConn(), "runtime", "workspace", proto.Version, registry, nil) + heartbeat, err := proto.NewEnvelope(proto.TypeHeartbeat, "", proto.HeartbeatPayload{SupportedAgentKinds: []proto.SupportedAgentKind{{Kind: "fixture", Available: true, Capabilities: prototest.Capabilities(proto.AgentKindCapabilities{})}}}) + if err != nil { + t.Fatal(err) + } + for _, want := range []bool{true, false} { + session.handleHeartbeat(heartbeat) + select { + case <-registry.Announced(): + if !want { + t.Fatal("repeated heartbeat announced again") + } + default: + if want { + t.Fatal("first agent kinds not announced") + } + } + } +} diff --git a/services/core/internal/runtimegateway/registry.go b/services/core/internal/runtimegateway/registry.go index 29ef0ddcf..45111df43 100644 --- a/services/core/internal/runtimegateway/registry.go +++ b/services/core/internal/runtimegateway/registry.go @@ -37,15 +37,31 @@ type Registry struct { // happens between WaitForDevice registering and selecting on the // chan still wakes the waiter. waiters map[string][]chan *Session + + // announced coalesces "a session reported its agent kinds for the first + // time" for the one owner that waits on known kinds. + announced chan struct{} } // NewRegistry returns an empty registry. The zero value would also // work but the constructor avoids accidental nil-map panics. func NewRegistry() *Registry { return &Registry{ - byDevice: map[string]*Session{}, - byRun: map[string]*Session{}, - waiters: map[string][]chan *Session{}, + byDevice: map[string]*Session{}, + byRun: map[string]*Session{}, + waiters: map[string][]chan *Session{}, + announced: make(chan struct{}, 1), + } +} + +// Announced signals after a session first reports its agent kinds. Polling +// recovers a coalesced signal. +func (r *Registry) Announced() <-chan struct{} { return r.announced } + +func (r *Registry) announce() { + select { + case r.announced <- struct{}{}: + default: } } diff --git a/services/core/internal/runtimegateway/session.go b/services/core/internal/runtimegateway/session.go index 1ec05993d..a2e48df53 100644 --- a/services/core/internal/runtimegateway/session.go +++ b/services/core/internal/runtimegateway/session.go @@ -209,13 +209,16 @@ func (s *Session) AgentKindStatus(kind string) (info runtimedevice.SupportedAgen return runtimedevice.SupportedAgentKind{}, false, true } -func (s *Session) setSupportedAgentKinds(kinds []runtimedevice.SupportedAgentKind) { +// setSupportedAgentKinds reports whether these are the session's first kinds. +func (s *Session) setSupportedAgentKinds(kinds []runtimedevice.SupportedAgentKind) (first bool) { copyKinds := make([]runtimedevice.SupportedAgentKind, len(kinds)) copy(copyKinds, kinds) s.kindsMu.Lock() + first = !s.kindsSeen s.kindsSeen = true s.supportedKinds = copyKinds s.kindsMu.Unlock() + return first } // Close closes the transport and subscriptions with ErrSessionClosed, then @@ -396,7 +399,9 @@ func (s *Session) handleHeartbeat(env proto.Envelope) { return } kinds := deviceKindsFromHeartbeat(p) - s.setSupportedAgentKinds(kinds) + if s.setSupportedAgentKinds(kinds) { + s.reg.announce() + } if s.heartbeat == nil { return }