admin/BishopFox-sliver
1
0
Fork 0
mirror of https://github.com/BishopFox/sliver synced 2026-06-08 10:29:05 +00:00
Code Issues Packages Projects Releases Wiki Activity

Wait for AI RPC event streams to subscribe

Browse Source
This commit is contained in:
moloch--
2026-04-02 09:17:33 -07:00
parent 3aefc39c08
commit c23067a0e0
1 changed files with 26 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
server/rpc/rpc-ai_test.go
+26 -4
View File
@@ -17,6 +17,7 @@ import (
"github.com/bishopfox/sliver/protobuf/rpcpb"
serverai "github.com/bishopfox/sliver/server/ai"
"github.com/bishopfox/sliver/server/configs"
"github.com/bishopfox/sliver/server/core"
"github.com/bishopfox/sliver/server/db"
"github.com/bishopfox/sliver/server/db/models"
"google.golang.org/grpc/codes"
@@ -73,7 +74,7 @@ func TestSaveAIConversationMessageCompletesConversationAndPublishesEvents(t *tes
streamCtx, cancelStream := context.WithTimeout(context.Background(), aiConversationEventStreamTimeout)
defer cancelStream()
eventStream, err := client.Events(streamCtx, &commonpb.Empty{})
eventStream, err := startAIEventStream(t, client, streamCtx)
if err != nil {
t.Fatalf("start events stream: %v", err)
}
@@ -246,7 +247,7 @@ func TestSaveAIConversationMessageCompletesOpenAIWithoutExplicitBaseURL(t *testi
streamCtx, cancelStream := context.WithTimeout(context.Background(), aiConversationEventStreamTimeout)
defer cancelStream()
eventStream, err := client.Events(streamCtx, &commonpb.Empty{})
eventStream, err := startAIEventStream(t, client, streamCtx)
if err != nil {
t.Fatalf("start events stream: %v", err)
}
@@ -336,7 +337,7 @@ func TestSaveAIConversationMessagePublishesFailureMessageWhenProviderErrors(t *t
streamCtx, cancelStream := context.WithTimeout(context.Background(), aiConversationEventStreamTimeout)
defer cancelStream()
eventStream, err := client.Events(streamCtx, &commonpb.Empty{})
eventStream, err := startAIEventStream(t, client, streamCtx)
if err != nil {
t.Fatalf("start events stream: %v", err)
}
@@ -626,7 +627,7 @@ func TestSaveAIConversationMessagePersistsReasoningAndToolBlocks(t *testing.T) {
streamCtx, cancelStream := context.WithTimeout(context.Background(), aiConversationEventStreamTimeout)
defer cancelStream()
eventStream, err := client.Events(streamCtx, &commonpb.Empty{})
eventStream, err := startAIEventStream(t, client, streamCtx)
if err != nil {
t.Fatalf("start events stream: %v", err)
}
@@ -935,6 +936,27 @@ func saveOpenAICompletionConfig(t *testing.T, model string, thinkingLevel string
}
}
func startAIEventStream(t *testing.T, client rpcpb.SliverRPCClient, streamCtx context.Context) (rpcpb.SliverRPC_EventsClient, error) {
t.Helper()
before := len(core.Clients.ActiveOperators())
eventStream, err := client.Events(streamCtx, &commonpb.Empty{})
if err != nil {
return nil, err
}
deadline := time.Now().Add(5 * time.Second)
for time.Now().Before(deadline) {
if len(core.Clients.ActiveOperators()) > before {
return eventStream, nil
}
time.Sleep(10 * time.Millisecond)
}
t.Fatalf("event stream did not subscribe before deadline")
return nil, nil
}
func waitForAIConversationEvent(t *testing.T, eventStream rpcpb.SliverRPC_EventsClient, conversationID string) {
t.Helper()