mirror of
https://github.com/Ed1s0nZ/CyberStrikeAI.git
synced 2026-08-11 05:30:42 +02:00
Add files via upload
This commit is contained in:
@@ -1185,6 +1185,8 @@ func (h *AgentHandler) createProgressCallback(runCtx context.Context, cancelRun
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
flushResponsePlan()
|
flushResponsePlan()
|
||||||
|
// 助手正文开始前,推理流通常已结束;落库以便刷新后「渗透测试详情」可回放
|
||||||
|
flushThinkingStreams()
|
||||||
respPlan.meta = nil
|
respPlan.meta = nil
|
||||||
if dataMap, ok := data.(map[string]interface{}); ok {
|
if dataMap, ok := data.(map[string]interface{}); ok {
|
||||||
respPlan.meta = make(map[string]interface{}, len(dataMap))
|
respPlan.meta = make(map[string]interface{}, len(dataMap))
|
||||||
@@ -1220,6 +1222,19 @@ func (h *AgentHandler) createProgressCallback(runCtx context.Context, cancelRun
|
|||||||
}
|
}
|
||||||
if eventType == "response" {
|
if eventType == "response" {
|
||||||
flushResponsePlan()
|
flushResponsePlan()
|
||||||
|
flushThinkingStreams()
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if eventType == "done" {
|
||||||
|
flushResponsePlan()
|
||||||
|
flushThinkingStreams()
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
// 流式思考/推理结束:聚合落库(与 eino_agent_reply_stream_end 同理)
|
||||||
|
if eventType == "thinking_stream_end" || eventType == "reasoning_chain_stream_end" {
|
||||||
|
flushResponsePlan()
|
||||||
|
flushThinkingStreams()
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -3,10 +3,14 @@ package handler
|
|||||||
import (
|
import (
|
||||||
"context"
|
"context"
|
||||||
"fmt"
|
"fmt"
|
||||||
|
"os"
|
||||||
|
"path/filepath"
|
||||||
"sync"
|
"sync"
|
||||||
"testing"
|
"testing"
|
||||||
|
|
||||||
"cyberstrike-ai/internal/config"
|
"cyberstrike-ai/internal/config"
|
||||||
|
"cyberstrike-ai/internal/database"
|
||||||
|
"cyberstrike-ai/internal/openai"
|
||||||
|
|
||||||
"go.uber.org/zap"
|
"go.uber.org/zap"
|
||||||
)
|
)
|
||||||
@@ -46,3 +50,50 @@ func TestCreateProgressCallback_ConcurrentToolEvents(t *testing.T) {
|
|||||||
}
|
}
|
||||||
wg.Wait()
|
wg.Wait()
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// TestCreateProgressCallback_FlushesReasoningOnDone 流式推理聚合须在 done/response 时落库,刷新后可回放。
|
||||||
|
func TestCreateProgressCallback_FlushesReasoningOnDone(t *testing.T) {
|
||||||
|
tmp := t.TempDir()
|
||||||
|
db, err := database.NewDB(filepath.Join(tmp, "test.sqlite"), zap.NewNop())
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("NewDB: %v", err)
|
||||||
|
}
|
||||||
|
defer os.RemoveAll(tmp)
|
||||||
|
|
||||||
|
conv, err := db.CreateConversation("test", database.ConversationCreateMeta{})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("CreateConversation: %v", err)
|
||||||
|
}
|
||||||
|
asst, err := db.AddMessage(conv.ID, "assistant", "处理中...", nil)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("AddMessage: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
h := &AgentHandler{logger: zap.NewNop(), db: db}
|
||||||
|
cb := h.createProgressCallback(context.Background(), nil, conv.ID, asst.ID, nil)
|
||||||
|
|
||||||
|
streamID := "eino-reasoning-test-1"
|
||||||
|
cb("reasoning_chain_stream_start", " ", map[string]interface{}{
|
||||||
|
"streamId": streamID,
|
||||||
|
"source": "eino",
|
||||||
|
})
|
||||||
|
cb("reasoning_chain_stream_delta", "step one", openai.WithSSEAccumulated(map[string]interface{}{
|
||||||
|
"streamId": streamID,
|
||||||
|
}, "step one"))
|
||||||
|
cb("done", "", map[string]interface{}{"conversationId": conv.ID})
|
||||||
|
|
||||||
|
details, err := db.GetProcessDetails(asst.ID)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("GetProcessDetails: %v", err)
|
||||||
|
}
|
||||||
|
found := false
|
||||||
|
for _, d := range details {
|
||||||
|
if d.EventType == "reasoning_chain" && d.Message == "step one" {
|
||||||
|
found = true
|
||||||
|
break
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if !found {
|
||||||
|
t.Fatalf("expected reasoning_chain persisted on done, got %+v", details)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user