From 101155bdc7a85f389bc697510ae9c7bec66c3b04 Mon Sep 17 00:00:00 2001 From: Haruko386 Date: Thu, 16 Jul 2026 19:08:36 +0800 Subject: [PATCH] fix: support message content (#16980) ### Summary As Title image --- internal/agent/canvas/runner.go | 26 +- internal/agent/component/agent.go | 116 ++++++++- internal/agent/component/agent_test.go | 85 +++++++ internal/agent/component/llm.go | 30 ++- internal/agent/component/message.go | 130 ++-------- .../agent/component/message_chunked_test.go | 85 ------- internal/agent/component/message_test.go | 19 +- internal/agent/runtime/context.go | 90 +++++++ internal/entity/models/llm.go | 18 +- internal/entity/models/types.go | 3 + internal/handler/agent.go | 27 ++- internal/handler/agent_test.go | 59 ++++- internal/service/agent.go | 195 +++++++++++++-- internal/service/agent_run_e2e_test.go | 9 +- internal/service/agent_test.go | 223 +++++++++++++++++- internal/service/bot_completion.go | 4 +- 16 files changed, 831 insertions(+), 288 deletions(-) delete mode 100644 internal/agent/component/message_chunked_test.go diff --git a/internal/agent/canvas/runner.go b/internal/agent/canvas/runner.go index 138fd7a997..003fa8fdb3 100644 --- a/internal/agent/canvas/runner.go +++ b/internal/agent/canvas/runner.go @@ -14,7 +14,7 @@ // limitations under the License. // -// runner.go — Canvas execution runtime. Drives a Canvas invocation +// Package canvas runner.go — Canvas execution runtime. Drives a Canvas invocation // (the caller supplies the RunFunc that does Compile+Invoke), catches // the four possible outcomes, and surfaces them as RunEvent values on // a channel that the HTTP layer streams as SSE frames. @@ -44,7 +44,6 @@ // - RunEvent.Type == "message" → {data: } // - RunEvent.Type == "waiting_for_user" → {cpn_id: } // - RunEvent.Type == "error" → {message: } -// - RunEvent.Type == "done" → final terminator frame package canvas import ( @@ -107,8 +106,11 @@ type NodeFinishedData struct { // MessageEvent is the JSON payload for Type=="message" frames. type MessageEvent struct { - Content string `json:"content"` - Reference interface{} `json:"reference,omitempty"` + Content string `json:"content"` + Reference interface{} `json:"reference,omitempty"` + Thinking string `json:"thinking,omitempty"` + StartToThink bool `json:"start_to_think,omitempty"` + EndToThink bool `json:"end_to_think,omitempty"` } // MessageEndEvent is the JSON payload for Type=="message_end" frames. @@ -335,11 +337,6 @@ func (r *Runner) Run( } } push(out, RunEvent{Type: "waiting_for_user", Data: safeEventJSON(waiting), MessageID: messageID, CreatedAt: nowUnix(), TaskID: taskID, SessionID: sessionID}) - // Always close a RunAgent call with the `done` - // terminator so the front-end can rely on a - // channel-end sentinel regardless of whether the run - // completed, errored, or paused for user input. - push(out, RunEvent{Type: "done", Data: "", MessageID: messageID, CreatedAt: nowUnix(), TaskID: taskID, SessionID: sessionID}) return } if IsInterruptError(runErr) { @@ -349,22 +346,11 @@ func (r *Runner) Run( // the first paused session it knows about. r.saveInterruptID(canvasID, sessionID, runErr.Error()) push(out, RunEvent{Type: "waiting_for_user", Data: safeEventJSON(WaitingForUserEvent{CpnID: runErr.Error()}), MessageID: messageID, CreatedAt: nowUnix(), TaskID: taskID, SessionID: sessionID}) - push(out, RunEvent{Type: "done", Data: "", MessageID: messageID, CreatedAt: nowUnix(), TaskID: taskID, SessionID: sessionID}) return } pushErr(out, runErr.Error()) - // Close the channel with the `done` terminator so the - // front-end sees a channel-end sentinel on the error - // path too — matches the contract for completed and - // waiting-for-user paths above. - push(out, RunEvent{Type: "done", Data: "", MessageID: messageID, CreatedAt: nowUnix(), TaskID: taskID, SessionID: sessionID}) return } - - // Normal completion — the buildRunFunc already emitted the - // workflow events during execution. Runner just sends the - // terminator. - push(out, RunEvent{Type: "done", Data: "", MessageID: messageID, CreatedAt: nowUnix(), TaskID: taskID, SessionID: sessionID}) }() return out diff --git a/internal/agent/component/agent.go b/internal/agent/component/agent.go index 28f47c8445..7278623df6 100644 --- a/internal/agent/component/agent.go +++ b/internal/agent/component/agent.go @@ -14,7 +14,9 @@ package component import ( "context" "encoding/json" + "errors" "fmt" + "io" "sort" "strings" @@ -72,6 +74,7 @@ type AgentParam struct { ModelID string SystemPrompt string UserPrompt string + Thinking string TopP *float64 Tools []string // Agent-visible tool names resolved into Eino BaseTool instances ToolParams map[string]map[string]any // node-level tool constructor params keyed by tool name @@ -135,7 +138,7 @@ var agentRunner = runEinoReActAgent // runEinoReActAgent creates an eino react agent and runs it against the // model built from p. func runEinoReActAgent(ctx context.Context, p AgentParam) (*schema.Message, error) { - chatModel, err := buildAgentChatModel(p) + chatModel, err := buildAgentChatModel(ctx, p) if err != nil { return nil, fmt.Errorf("build model: %w", err) } @@ -164,13 +167,89 @@ func runEinoReActAgent(ctx context.Context, p AgentParam) (*schema.Message, erro input := []*schema.Message{schema.UserMessage(p.UserPrompt)} opt, future := react.WithMessageFuture() ctx = setArtifactCollector(ctx, future) - msg, err := agent.Generate(ctx, input, opt) + stream, err := agent.Stream(ctx, input, opt) + if err != nil { + return nil, err + } + defer stream.Close() + emitDone := emitAgentModelStreams(ctx, future) + + chunks := make([]*schema.Message, 0) + for { + chunk, err := stream.Recv() + if err == io.EOF { + break + } + if err != nil { + return nil, err + } + if chunk == nil { + continue + } + chunks = append(chunks, chunk) + } + if emitErr := <-emitDone; emitErr != nil { + return nil, emitErr + } + if len(chunks) == 0 { + return &schema.Message{Role: schema.Assistant}, nil + } + msg, err := schema.ConcatMessages(chunks) if err != nil { return nil, err } return msg, nil } +func emitAgentModelStreams(ctx context.Context, future react.MessageFuture) <-chan error { + done := make(chan error, 1) + go func() { + var firstErr error + iter := future.GetMessageStreams() + for { + msgStream, hasNext, err := iter.Next() + if err != nil { + firstErr = err + break + } + if !hasNext { + break + } + if msgStream == nil { + continue + } + for { + msg, err := msgStream.Recv() + if errors.Is(err, io.EOF) { + break + } + if err != nil { + if firstErr == nil { + firstErr = err + } + break + } + if msg == nil { + continue + } + if msg.Role != "" && msg.Role != schema.Assistant { + continue + } + if msg.Content == "" && msg.ReasoningContent == "" { + continue + } + if runtime.AgentMessageEventsEmitted(ctx) { + continue + } + runtime.EmitAgentMessage(ctx, msg.Content, msg.ReasoningContent) + } + msgStream.Close() + } + done <- firstErr + }() + return done +} + // addToolCallMemory summarizes the tool calls observed in msg via // a small LLM call and returns a one-line history entry. Mirrors // Python's `add_memory(user, assist, func_name, params, results, @@ -437,6 +516,9 @@ func (c *AgentComponent) Name() string { return "Agent" } // Invoke runs the ReAct loop via the configured agentRunner and returns // the output map. func (c *AgentComponent) Invoke(ctx context.Context, inputs map[string]any) (map[string]any, error) { + runtime.ResetAgentMessageEmission(ctx) + defer runtime.FinalizeAgentMessage(ctx) + p := mergeAgentParam(c.param, inputs) hasRuntimeUserPrompt := false if v, ok := stringFrom(inputs, "user_prompt"); ok { @@ -540,6 +622,8 @@ func (c *AgentComponent) Invoke(ctx context.Context, inputs map[string]any) (map zap.String("modelID", p.ModelID), zap.Int("content_len", len(msg.Content))) content := msg.Content + thinking := msg.ReasoningContent + var groundingStatus string if p.Cite { chunks := chunksFromState(ctx) @@ -559,12 +643,16 @@ func (c *AgentComponent) Invoke(ctx context.Context, inputs map[string]any) (map artifactMD := formatArtifactMarkdown(artifacts, content) out := map[string]any{ "content": content + artifactMD, + "thinking": thinking, "tool_calls": extractToolCalls(msg), "artifacts": artifacts, } if groundingStatus != "" { out["grounding_status"] = groundingStatus } + if !runtime.AgentMessageEventsEmitted(ctx) { + runtime.EmitAgentMessage(ctx, content+artifactMD, thinking) + } return out, nil } @@ -607,6 +695,7 @@ func (c *AgentComponent) Inputs() map[string]string { func (c *AgentComponent) Outputs() map[string]string { return map[string]string{ "content": "Final assistant content (after the ReAct loop terminates)", + "thinking": "Model reasoning content, when the provider returns it separately.", "tool_calls": "One entry per tool call observed during the run", "artifacts": "Artifacts collected from tool responses (empty in P0)", "grounding_status": "'applied' | 'no_chunks' | 'error: ' (present when cite=true).", @@ -615,7 +704,7 @@ func (c *AgentComponent) Outputs() map[string]string { // buildAgentChatModel constructs an EinoChatModel from AgentParam by // resolving the driver through the RAGFlow provider manager. -func buildAgentChatModel(p AgentParam) (*models.EinoChatModel, error) { +func buildAgentChatModel(ctx context.Context, p AgentParam) (*models.EinoChatModel, error) { driver := p.Driver modelID := p.ModelID @@ -665,8 +754,21 @@ func buildAgentChatModel(p AgentParam) (*models.EinoChatModel, error) { // would be dead weight. When AgentParam grows Temperature/ // MaxTokens, switch to always-build. var chatCfg *models.ChatConfig - if p.TopP != nil { + if p.TopP != nil || p.Thinking != "" || runtime.HasAgentMessageEmitter(ctx) { chatCfg = &models.ChatConfig{TopP: p.TopP} + switch p.Thinking { + case "enabled": + t := true + chatCfg.Thinking = &t + case "disabled": + f := false + chatCfg.Thinking = &f + } + if runtime.HasAgentMessageEmitter(ctx) { + chatCfg.StreamCallback = func(contentDelta, reasoningDelta string) { + runtime.EmitAgentMessage(ctx, contentDelta, reasoningDelta) + } + } } return models.NewEinoChatModel(cm, chatCfg), nil } @@ -976,6 +1078,9 @@ func mergeAgentParam(base AgentParam, inputs map[string]any) AgentParam { f := v p.TopP = &f } + if v, ok := stringFrom(inputs, "thinking"); ok && v != "" && v != "default" { + p.Thinking = v + } if v, ok := intFrom(inputs, "max_rounds"); ok { p.MaxRounds = v } @@ -1182,6 +1287,9 @@ func init() { if v, ok := nestedMapFrom(params, "tool_params"); ok { p.ToolParams = mergeToolParams(p.ToolParams, v) } + if v, ok := stringFrom(params, "thinking"); ok && v != "" && v != "default" { + p.Thinking = v + } if v, ok := intFrom(params, "max_rounds"); ok { p.MaxRounds = v } diff --git a/internal/agent/component/agent_test.go b/internal/agent/component/agent_test.go index 63d9dd5a19..a21c472c75 100644 --- a/internal/agent/component/agent_test.go +++ b/internal/agent/component/agent_test.go @@ -70,6 +70,91 @@ func TestAgent_NoToolsReAct(t *testing.T) { } } +func TestAgent_EmitsThinking(t *testing.T) { + withAgentRunner(t, func(_ context.Context, _ AgentParam) (*schema.Message, error) { + return &schema.Message{ + Role: schema.Assistant, + Content: "final answer", + ReasoningContent: "model reasoning", + }, nil + }) + + c := NewAgentComponent(AgentParam{ModelID: "stub", MaxRounds: 1}) + out, err := c.Invoke(context.Background(), map[string]any{ + "user_prompt": "hello", + }) + if err != nil { + t.Fatalf("Invoke: %v", err) + } + if got, want := out["content"], "final answer"; got != want { + t.Errorf("content=%v, want %v", got, want) + } + if got, want := out["thinking"], "model reasoning"; got != want { + t.Errorf("thinking=%v, want %v", got, want) + } +} + +func TestAgent_MessageEmissionIsScopedPerInvocation(t *testing.T) { + responses := []string{"first answer", "second answer"} + withAgentRunner(t, func(_ context.Context, _ AgentParam) (*schema.Message, error) { + if len(responses) == 0 { + t.Fatal("agent runner called too many times") + } + content := responses[0] + responses = responses[1:] + return &schema.Message{Role: schema.Assistant, Content: content}, nil + }) + + state := runtime.NewCanvasState("run-1", "task-1") + ctx := runtime.WithState(context.Background(), state) + var contents []string + ctx = runtime.WithAgentMessageEmitterControl(ctx, + func(contentDelta, _ string) { + if contentDelta != "" { + contents = append(contents, contentDelta) + } + }, + func() bool { return false }, + func() {}, + ) + + first := NewAgentComponent(AgentParam{ModelID: "stub", MaxRounds: 1}) + if _, err := first.Invoke(ctx, map[string]any{"user_prompt": "first"}); err != nil { + t.Fatalf("first Invoke: %v", err) + } + second := NewAgentComponent(AgentParam{ModelID: "stub", MaxRounds: 1}) + if _, err := second.Invoke(ctx, map[string]any{"user_prompt": "second"}); err != nil { + t.Fatalf("second Invoke: %v", err) + } + + if got, want := strings.Join(contents, "|"), "first answer|second answer"; got != want { + t.Fatalf("emitted contents = %q, want %q", got, want) + } +} + +func TestAgent_ForwardsThinkingParam(t *testing.T) { + var gotThinking string + withAgentRunner(t, func(_ context.Context, p AgentParam) (*schema.Message, error) { + gotThinking = p.Thinking + return &schema.Message{Role: schema.Assistant, Content: "ok"}, nil + }) + + cmp, err := New("Agent", map[string]any{ + "model_id": "stub", + "user_prompt": "hello", + "thinking": "enabled", + }) + if err != nil { + t.Fatalf("New(Agent): %v", err) + } + if _, err := cmp.Invoke(context.Background(), nil); err != nil { + t.Fatalf("Invoke: %v", err) + } + if gotThinking != "enabled" { + t.Fatalf("runner thinking = %q, want enabled", gotThinking) + } +} + func TestAgent_ResolvesUserPromptFromCanvasState(t *testing.T) { var gotPrompt string withAgentRunner(t, func(_ context.Context, p AgentParam) (*schema.Message, error) { diff --git a/internal/agent/component/llm.go b/internal/agent/component/llm.go index 3370110efc..09c1b90423 100644 --- a/internal/agent/component/llm.go +++ b/internal/agent/component/llm.go @@ -162,10 +162,11 @@ type ChatInvokeRequest struct { // ChatInvokeResponse mirrors what the LLM component writes to its outputs. type ChatInvokeResponse struct { - Content string - Model string - Stopped bool - Tokens int + Content string + Thinking string + Model string + Stopped bool + Tokens int } // defaultChatInvokerMu guards defaultChatInvoker swaps during tests. @@ -262,10 +263,11 @@ func (e *einoChatInvoker) Invoke(ctx context.Context, req ChatInvokeRequest) (*C return nil, err } return &ChatInvokeResponse{ - Content: out.Content, - Model: modelName, - Stopped: true, - Tokens: 0, + Content: out.Content, + Thinking: out.ReasoningContent, + Model: modelName, + Stopped: true, + Tokens: 0, }, nil } @@ -552,10 +554,11 @@ func (c *LLMComponent) Invoke(ctx context.Context, inputs map[string]any) (map[s } out := map[string]any{ - "content": cleaned, - "model": resp.Model, - "stopped": resp.Stopped, - "tokens": resp.Tokens, + "content": cleaned, + "thinking": resp.Thinking, + "model": resp.Model, + "stopped": resp.Stopped, + "tokens": resp.Tokens, } if p.JSONOutput { var parsed map[string]any @@ -603,6 +606,7 @@ func (c *LLMComponent) Invoke(ctx context.Context, inputs map[string]any) (map[s common.Warn("component: LLM: output_structure set but no parseable JSON after retry") } } + out["thinking"] = resp.Thinking return out, nil } @@ -651,7 +655,7 @@ func (c *LLMComponent) Stream(ctx context.Context, inputs map[string]any) (<-cha // A real streaming integration would loop over a channel // here and emit multiple chunks with partial content. chunk := map[string]any{ - "thinking": "", + "thinking": result["thinking"], "content": result["content"], } select { diff --git a/internal/agent/component/message.go b/internal/agent/component/message.go index b4544e864e..3593e421dc 100644 --- a/internal/agent/component/message.go +++ b/internal/agent/component/message.go @@ -33,8 +33,6 @@ package component import ( "context" "fmt" - "maps" - "regexp" "strings" "ragflow/internal/agent/audio" @@ -147,9 +145,7 @@ func (m *MessageComponent) Name() string { return m.name } // Invoke resolves inputs["text"] (or the per-instance text seeded // from params at build time) as a template against the current -// *CanvasState, returns the resolved string at outputs["content"], and -// (if inputs["stream"] == true) records the number of chunks in -// outputs["streamed_chunks"]. +// *CanvasState and returns the resolved string at outputs["content"]. // // Message Invoke behaviour: // - input-format override: inputs["output_format"] wins over the @@ -183,14 +179,11 @@ func (m *MessageComponent) Invoke(ctx context.Context, inputs map[string]any) (m if text == "" { text = fallbackMessageText(inputs) } - // Message is a display node, not parameter binding. Use the - // tolerant resolver (nil refs render as empty string) instead - // of runtime.ResolveTemplate — matches the Python canvas.py - // soft-fail semantic so authoring patterns like - // {Component@head} for optional fields don't crash the run when - // the upstream list is empty. Parameter-binding call sites keep - // the loud-fail contract via runtime.ResolveTemplate. - resolved := runtime.ResolveTemplateForDisplay(text, state) + + resolved, err := runtime.ResolveTemplate(text, state) + if err != nil { + return nil, fmt.Errorf("Message: %w", err) + } // Extract downloads. Walks inputs for download-info maps so // callers can attach binaries to the message body. @@ -278,12 +271,6 @@ func (m *MessageComponent) Invoke(ctx context.Context, inputs map[string]any) (m memSave, _ := inputs["memory_save"].(bool) if memSave { memIDs := extractMemoryIDs(inputs) - if len(memIDs) == 0 { - // Fall back to per-instance memory_ids declared in - // the DSL — the orchestrator may not re-pass them - // when it overrides only `memory_save`. - memIDs = extractMemoryIDsFromParams(m.text) - } if len(memIDs) > 0 { saver := GetMemorySaver() saveErr := saver.Save(ctx, MemorySaveRequest{ @@ -300,11 +287,6 @@ func (m *MessageComponent) Invoke(ctx context.Context, inputs map[string]any) (m } } - if streamOn, _ := inputs["stream"].(bool); streamOn { - // P0: one chunk for the whole resolved content. A later phase - // can split on token / sentence boundaries. - out["streamed_chunks"] = 1 - } return out, nil } @@ -370,15 +352,6 @@ func isMessageInfraInput(key string) bool { } } -// extractMemoryIDsFromParams looks for a "_memory_ids" hint in -// the component's stored text — used as a last-ditch fallback -// when the orchestrator does not re-pass memory_ids. Returns nil -// in the common case; this helper exists to keep the public -// memory-save flow permissive about caller omissions. -func extractMemoryIDsFromParams(_ string) []string { - return nil -} - // stringFromStateSys reads a sys-level state value. Returns "" // when state or the key is missing. Used by the memory-save path // to pull the user's original query. @@ -394,17 +367,9 @@ func stringFromStateSys(state *runtime.CanvasState, key string) string { return "" } -// Stream is the SSE variant. The resolved template content is -// split on sentence boundaries ([.!?]\s+ between letters/digits) -// and each sentence is emitted as a separate chunk. A trailing -// "done" marker signals end-of-stream. The chunk map's "content" -// key carries the sentence text; "done" is true on the final -// chunk. -// -// The splitter uses a regex for portable sentence boundaries -// without pulling in a tokenizer. A future tokenizer-aware -// splitter (gonja + langdetect, or a small Go -// sentence-segmentation lib) can improve break quality. +// Stream resolves the message and emits the content chunk. The outer +// Agent SSE handler owns the final [DONE] frame, matching Python's +// agent_api.py rather than leaking a component-local done marker. func (m *MessageComponent) Stream(ctx context.Context, inputs map[string]any) (<-chan map[string]any, error) { ch := make(chan map[string]any, 16) go func() { @@ -418,67 +383,14 @@ func (m *MessageComponent) Stream(ctx context.Context, inputs map[string]any) (< return } text, _ := result["content"].(string) - if text == "" { - // Nothing to split; emit a single empty-content chunk - // plus the done marker so downstream consumers have a - // well-defined two-chunk stream. - select { - case ch <- map[string]any{"content": "", "thinking": ""}: - case <-ctx.Done(): - return - } - } else { - sentences := splitSentences(text) - for _, s := range sentences { - select { - case ch <- map[string]any{"content": s, "thinking": ""}: - case <-ctx.Done(): - return - } - } - } select { - case ch <- map[string]any{"done": true, "model": result["model"]}: + case ch <- map[string]any{"content": text, "thinking": ""}: case <-ctx.Done(): } }() return ch, nil } -// sentenceSplitRe matches sentence boundaries: ".", "!", or "?" -// followed by whitespace. The character class keeps the -// abbreviations list short for v1; the follow-up tokenizer-aware -// splitter is a more robust replacement. -var sentenceSplitRe = regexp.MustCompile(`([.!?])\s+`) - -// splitSentences splits text on sentence boundaries, preserving -// the trailing punctuation. Returns a slice of at least one -// element; empty input returns a single empty element. -func splitSentences(text string) []string { - if text == "" { - return []string{""} - } - matches := sentenceSplitRe.FindAllStringIndex(text, -1) - if len(matches) == 0 { - return []string{text} - } - out := make([]string, 0, len(matches)+1) - prev := 0 - for _, m := range matches { - // Include the matched punctuation in the previous sentence - // but stop BEFORE the trailing whitespace — otherwise each - // emitted sentence has a leading space, which both the - // Message component's stream joiner and the v1 Python - // chunker would have to re-trim. - out = append(out, text[prev:m[0]+1]) - prev = m[1] - } - if prev < len(text) { - out = append(out, text[prev:]) - } - return out -} - // Inputs returns the public parameter surface. Field types match // the Python DSL contract (text template, stream toggle, // memory_save toggle). @@ -495,27 +407,17 @@ func (m *MessageComponent) Inputs() map[string]string { } } -// Outputs returns the resolved template plus the streamed-chunk -// counter. +// Outputs returns the resolved template plus optional side-channel outputs. func (m *MessageComponent) Outputs() map[string]string { return map[string]string{ - "content": "Resolved and rendered message body.", - "streamed_chunks": "Number of SSE chunks emitted (present when stream=true).", - "downloads": "Extracted download descriptors ({doc_id, filename, mime_type, url}).", - "audio": "{media_type, data_b64} envelope populated when auto_play is wired and a TTS engine succeeds.", - "audio_error": "Surfaced when TTS dispatch fails; the textual content is still returned.", - "memory_error": "Surfaced when memory persistence fails; the textual content is still returned.", + "content": "Resolved and rendered message body.", + "downloads": "Extracted download descriptors ({doc_id, filename, mime_type, url}).", + "audio": "{media_type, data_b64} envelope populated when auto_play is wired and a TTS engine succeeds.", + "audio_error": "Surfaced when TTS dispatch fails; the textual content is still returned.", + "memory_error": "Surfaced when memory persistence fails; the textual content is still returned.", } } -// mapCopy shallow-copies src into a fresh map. Used to keep Message's -// passthrough outputs un-aliased from the caller's inputs map. -func mapCopy(src map[string]any) map[string]any { - out := make(map[string]any, len(src)) - maps.Copy(out, src) - return out -} - func init() { Register(componentNameMessage, NewMessageComponent) } diff --git a/internal/agent/component/message_chunked_test.go b/internal/agent/component/message_chunked_test.go deleted file mode 100644 index 2a97088004..0000000000 --- a/internal/agent/component/message_chunked_test.go +++ /dev/null @@ -1,85 +0,0 @@ -// -// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. -// - -package component - -import ( - "strings" - "testing" -) - -// TestSplitSentences_Empty: empty input returns a single empty -// element (never nil, callers can always range). -func TestSplitSentences_Empty(t *testing.T) { - out := splitSentences("") - if len(out) != 1 { - t.Fatalf("expected 1 element, got %d", len(out)) - } - if out[0] != "" { - t.Errorf("expected empty string, got %q", out[0]) - } -} - -// TestSplitSentences_NoBoundary: text without sentence boundaries -// returns the whole text as a single element. -func TestSplitSentences_NoBoundary(t *testing.T) { - out := splitSentences("just a fragment") - if len(out) != 1 { - t.Fatalf("expected 1 element, got %d", len(out)) - } - if out[0] != "just a fragment" { - t.Errorf("got %q, want full text", out[0]) - } -} - -// TestSplitSentences_DotBoundary: a period+space split. -func TestSplitSentences_DotBoundary(t *testing.T) { - out := splitSentences("first. second. third.") - if len(out) != 3 { - t.Fatalf("expected 3 sentences, got %d: %+v", len(out), out) - } - if out[0] != "first." || out[1] != "second." || out[2] != "third." { - t.Errorf("unexpected split: %+v", out) - } -} - -// TestSplitSentences_MixedPunctuation: "!" and "?" are also -// sentence boundaries. -func TestSplitSentences_MixedPunctuation(t *testing.T) { - out := splitSentences("wow! really? ok.") - if len(out) != 3 { - t.Fatalf("expected 3, got %d: %+v", len(out), out) - } - if !strings.HasSuffix(out[0], "!") || !strings.HasSuffix(out[1], "?") || !strings.HasSuffix(out[2], ".") { - t.Errorf("wrong trailing punctuation: %+v", out) - } -} - -// TestMessage_Stream_MultiSentence: a resolved content with -// multiple sentences emits N content chunks + 1 done chunk. -// The end-to-end message_test.go already exercises the single- -// sentence happy path; this test pins the multi-sentence path -// through splitSentences directly (Stream integration is via the -// existing test). -func TestMessage_Stream_MultiSentence(t *testing.T) { - // Direct unit test of splitSentences; the Stream() method uses - // this internally. The full Stream() integration is covered by - // TestMessage_Stream in message_test.go. - sentences := splitSentences("First sentence. Second sentence. Third.") - if len(sentences) != 3 { - t.Fatalf("expected 3 sentences, got %d: %+v", len(sentences), sentences) - } -} diff --git a/internal/agent/component/message_test.go b/internal/agent/component/message_test.go index 19a659d24a..33d7d3e29e 100644 --- a/internal/agent/component/message_test.go +++ b/internal/agent/component/message_test.go @@ -25,8 +25,7 @@ import ( // TestMessage_ResolveTemplate asserts the canonical {{sys.x}} substitution // flow: a state with sys.query="world" and a template "hello {{sys.query}}" -// must resolve to "hello world". streamed_chunks must NOT be set when -// stream=false. +// must resolve to "hello world". func TestMessage_ResolveTemplate(t *testing.T) { c, _ := NewMessageComponent(nil) state := canvas.NewCanvasState("run-1", "task-1") @@ -45,13 +44,13 @@ func TestMessage_ResolveTemplate(t *testing.T) { t.Errorf("content: got %q, want %q", got, "hello world") } if _, ok := out["streamed_chunks"]; ok { - t.Errorf("streamed_chunks must not be present when stream=false, got %v", out["streamed_chunks"]) + t.Errorf("streamed_chunks must not be present, got %v", out["streamed_chunks"]) } } // TestMessage_Stream confirms the Stream() contract: the returned -// channel receives exactly one payload (the resolved content + -// streamed_chunks=1) and then closes. +// channel receives the resolved content and then closes. The outer +// SSE handler owns the [DONE] terminator. func TestMessage_Stream(t *testing.T) { c, _ := NewMessageComponent(nil) state := canvas.NewCanvasState("run-2", "task-2") @@ -65,20 +64,18 @@ func TestMessage_Stream(t *testing.T) { if err != nil { t.Fatalf("Stream: %v", err) } - // Chunked streaming. "hi" has no sentence boundary, so we - // expect exactly one content chunk plus the done marker. var got []map[string]any for chunk := range ch { got = append(got, chunk) } - if len(got) != 2 { - t.Fatalf("expected 2 chunks (content + done), got %d: %+v", len(got), got) + if len(got) != 1 { + t.Fatalf("expected 1 content chunk, got %d: %+v", len(got), got) } if got[0]["content"] != "hi" { t.Errorf("chunk[0].content=%q, want 'hi'", got[0]["content"]) } - if got[1]["done"] != true { - t.Errorf("chunk[1].done=%v, want true", got[1]["done"]) + if _, ok := got[0]["done"]; ok { + t.Errorf("component stream must not emit done marker: %+v", got[0]) } } diff --git a/internal/agent/runtime/context.go b/internal/agent/runtime/context.go index b9efbce4d2..af288ce68e 100644 --- a/internal/agent/runtime/context.go +++ b/internal/agent/runtime/context.go @@ -36,6 +36,19 @@ import ( // stable across calls (a fresh struct{}{} per call would key // distinctly and break ctx.Value lookups). type stateCtxKey struct{} +type agentMessageEmitterCtxKey struct{} + +// AgentMessageEmitter emits visible assistant deltas for an Agent component. +// The service layer owns the actual SSE envelope; runtime keeps the callback +// shape free of canvas/service imports. +type AgentMessageEmitter func(contentDelta, thinkingDelta string) + +type agentMessageEmitterState struct { + emit AgentMessageEmitter + finalize func() bool + reset func() + emitted bool +} // WithState attaches *CanvasState to ctx for retrieval by // GetStateFromContext. Production code (canvas/compile.go) calls this @@ -45,6 +58,83 @@ func WithState(ctx context.Context, s *CanvasState) context.Context { return context.WithValue(ctx, stateCtxKey{}, s) } +// WithAgentMessageEmitter attaches the Agent message stream callback used by +// components that can surface thinking before their node_finished event. +func WithAgentMessageEmitter(ctx context.Context, emit AgentMessageEmitter, finalize ...func() bool) context.Context { + if emit == nil { + return ctx + } + state := &agentMessageEmitterState{emit: emit} + if len(finalize) > 0 { + state.finalize = finalize[0] + } + return context.WithValue(ctx, agentMessageEmitterCtxKey{}, state) +} + +// WithAgentMessageEmitterControl attaches the Agent message stream callback +// with explicit lifecycle hooks for invocation-scoped reset/finalization. +func WithAgentMessageEmitterControl(ctx context.Context, emit AgentMessageEmitter, finalize func() bool, reset func()) context.Context { + if emit == nil { + return ctx + } + return context.WithValue(ctx, agentMessageEmitterCtxKey{}, &agentMessageEmitterState{ + emit: emit, + finalize: finalize, + reset: reset, + }) +} + +// HasAgentMessageEmitter reports whether the service layer installed an +// Agent message stream callback on ctx. +func HasAgentMessageEmitter(ctx context.Context) bool { + state, ok := ctx.Value(agentMessageEmitterCtxKey{}).(*agentMessageEmitterState) + return ok && state != nil && state.emit != nil +} + +// EmitAgentMessage emits Agent answer/thinking deltas when the service layer +// installed a callback. It returns true when a callback was present. +func EmitAgentMessage(ctx context.Context, contentDelta, thinkingDelta string) bool { + state, ok := ctx.Value(agentMessageEmitterCtxKey{}).(*agentMessageEmitterState) + if !ok || state == nil || state.emit == nil { + return false + } + state.emit(contentDelta, thinkingDelta) + if contentDelta != "" || thinkingDelta != "" { + state.emitted = true + } + return true +} + +// AgentMessageEventsEmitted reports whether the invocation-scoped Agent +// message emitter has emitted any deltas during the current run. +func AgentMessageEventsEmitted(ctx context.Context) bool { + state, ok := ctx.Value(agentMessageEmitterCtxKey{}).(*agentMessageEmitterState) + return ok && state != nil && state.emitted +} + +// FinalizeAgentMessage flushes the invocation-scoped Agent message emitter. +func FinalizeAgentMessage(ctx context.Context) { + state, ok := ctx.Value(agentMessageEmitterCtxKey{}).(*agentMessageEmitterState) + if !ok || state == nil || state.finalize == nil { + return + } + if state.finalize() { + state.emitted = true + } +} + +// ResetAgentMessageEmission starts a fresh Agent message emission scope. +func ResetAgentMessageEmission(ctx context.Context) { + state, ok := ctx.Value(agentMessageEmitterCtxKey{}).(*agentMessageEmitterState) + if !ok || state == nil { + return + } + if state.reset != nil { + state.reset() + } + state.emitted = false +} + // GetStateFromContext extracts a typed state attached via WithState. // Returns the state and a nil *sync.Mutex for *CanvasState (the // embedded RWMutex is what callers actually contend on through diff --git a/internal/entity/models/llm.go b/internal/entity/models/llm.go index 1e5fdd0bf4..5f0e990a7c 100644 --- a/internal/entity/models/llm.go +++ b/internal/entity/models/llm.go @@ -287,15 +287,23 @@ func (m *EinoChatModel) Stream(ctx context.Context, msgs []*schema.Message, opts sr, sw := schema.Pipe[*schema.Message](1) var sendMu sync.Mutex - sender := func(content *string, _ *string) error { + sender := func(content *string, reasoning *string) error { sendMu.Lock() defer sendMu.Unlock() - if content == nil { + if content == nil && reasoning == nil { return nil } - // Copy the string — the underlying buffer may be reused. - chunk := *content - if closed := sw.Send(&schema.Message{Role: schema.Assistant, Content: chunk}, nil); closed { + msg := &schema.Message{Role: schema.Assistant} + if content != nil { + msg.Content = *content + } + if reasoning != nil { + msg.ReasoningContent = *reasoning + } + if m.chatCfg != nil && m.chatCfg.StreamCallback != nil { + m.chatCfg.StreamCallback(msg.Content, msg.ReasoningContent) + } + if closed := sw.Send(msg, nil); closed { return fmt.Errorf("models: stream closed before send completed") } return nil diff --git a/internal/entity/models/types.go b/internal/entity/models/types.go index 0926075c47..38b56a2301 100644 --- a/internal/entity/models/types.go +++ b/internal/entity/models/types.go @@ -172,6 +172,9 @@ type ChatConfig struct { // non-nil) after the stream completes; callers read it the same // way they read ToolCallsResult. UsageResult *ChatUsage `json:"-"` + // StreamCallback receives raw content/reasoning deltas as soon as + // the model driver streams them. + StreamCallback func(contentDelta, reasoningDelta string) `json:"-"` } type APIConfig struct { diff --git a/internal/handler/agent.go b/internal/handler/agent.go index 611ae44799..59945ca51f 100644 --- a/internal/handler/agent.go +++ b/internal/handler/agent.go @@ -402,7 +402,11 @@ func (h *AgentHandler) RunAgent(c *gin.Context) { c.Writer.Header().Set("Content-Type", "text/event-stream") c.Writer.Header().Set("Cache-Control", "no-cache") c.Writer.Header().Set("Connection", "keep-alive") + doneSent := false for ev := range events { + if ev.Type == "done" { + doneSent = true + } if err := service.WriteChatbotRunEvent(c.Writer, ev); err != nil { common.Debug("agent run: client disconnected", zap.String("canvas_id", canvasID), @@ -412,6 +416,15 @@ func (h *AgentHandler) RunAgent(c *gin.Context) { return } } + if !doneSent { + if err := service.WriteChatbotRunEvent(c.Writer, canvas.RunEvent{Type: "done"}); err != nil { + common.Debug("agent run: failed to write [DONE]", + zap.String("canvas_id", canvasID), + zap.String("session_id", sessionID), + zap.Error(err), + ) + } + } } // readUserInput extracts the user_input field from the JSON body if @@ -813,7 +826,7 @@ func (h *AgentHandler) DeleteAgentSession(c *gin.Context) { // api/apps/restful_apis/agent_api.py:1440-1676. // // - When `stream` is true: streams SSE — one `data: {...}\n\n` frame per -// canvas RunEvent, terminated by `data: [DONE]\n\n`. +// canvas RunEvent, terminated by `data:[DONE]\n\n`. // - When `stream` is omitted or false (default, matches the Python // agent_chat_completion contract where `req.get("stream", False)` // defaults to non-streaming): collects all canvas events and returns a @@ -1008,7 +1021,7 @@ func (h *AgentHandler) AgentChatCompletions(c *gin.Context) { if req.Stream { // SSE streaming: one `data: {...}\n\n` frame per canvas RunEvent, - // terminated by `data: [DONE]\n\n`. We do NOT emit an SSE `event:` + // terminated by `data:[DONE]\n\n`. We do NOT emit an SSE `event:` // line — the front-end's use-send-message.ts parser feeds each // `data:` line directly into JSON.parse and expects the event type // in the JSON object's top-level `event` field. @@ -1016,8 +1029,12 @@ func (h *AgentHandler) AgentChatCompletions(c *gin.Context) { c.Writer.Header().Set("Cache-Control", "no-cache") c.Writer.Header().Set("Connection", "keep-alive") emitted := false + doneSent := false for ev := range events { emitted = true + if ev.Type == "done" { + doneSent = true + } common.Debug("agent chat completions: streaming event", zap.String("agent_id", req.AgentID), zap.String("session_id", req.SessionID), @@ -1037,7 +1054,7 @@ func (h *AgentHandler) AgentChatCompletions(c *gin.Context) { // Canvas produced no events (e.g. empty query). Echo the // session_id so the client can resume the conversation // (fixes #15169). The [DONE] terminator must be emitted - // here explicitly because the canvas never sends a + // after this branch because the canvas never sends a // "done" event on this path. common.Info("empty agent output - returning session_id", zap.String("agent_id", req.AgentID), @@ -1050,7 +1067,9 @@ func (h *AgentHandler) AgentChatCompletions(c *gin.Context) { SessionID: req.SessionID, } _ = service.WriteChatbotRunEvent(c.Writer, event) - if _, err := c.Writer.Write([]byte("data: [DONE]\n\n")); err != nil { + } + if !doneSent { + if err := service.WriteChatbotRunEvent(c.Writer, canvas.RunEvent{Type: "done"}); err != nil { common.Debug("agent chat completions: failed to write [DONE]", zap.Error(err), ) diff --git a/internal/handler/agent_test.go b/internal/handler/agent_test.go index 77cbdc52be..e98e316c29 100644 --- a/internal/handler/agent_test.go +++ b/internal/handler/agent_test.go @@ -708,7 +708,7 @@ func TestAgentChatCompletions_OpenAICompat_EmptyMessages(t *testing.T) { // SSE tests. It emits a pre-configured sequence of canvas.RunEvent // values on its RunAgent channel and then closes — enough to verify // the SSE wire format (Content-Type, one `data: {...}\n\n` frame per -// event, trailing `data: [DONE]\n\n`) without standing up the eino +// event, trailing `data:[DONE]\n\n`) without standing up the eino // runner or a live DB. type stubChatRunner struct { events []canvas.RunEvent @@ -729,7 +729,7 @@ func (s *stubChatRunner) RunAgent(_ context.Context, _, _, _, _ string, _ any) ( // TestAgentChatCompletions_StreamSetsContentType covers the SSE // path: the handler streams canvas.RunEvent frames as -// `data: {...}\n\n` with a trailing `data: [DONE]\n\n` terminator. +// `data: {...}\n\n` with a trailing `data:[DONE]\n\n` terminator. // The frame shape is the Python agent-canvas envelope // {event,message_id,task_id,session_id,data:{content}}. See // service.WriteChatbotRunEvent. @@ -765,7 +765,58 @@ func TestAgentChatCompletions_StreamSetsContentType(t *testing.T) { !strings.Contains(body, `"content":"hi back"`) { t.Errorf("body should contain flat agent event with content, got %q", body) } - if !strings.HasSuffix(body, "data: [DONE]\n\n") { + if !strings.HasSuffix(body, "data:[DONE]\n\n") { + t.Errorf("body should end with [DONE] terminator, got %q", body) + } +} + +func TestAgentChatCompletions_StreamAddsDoneWhenRunnerCloses(t *testing.T) { + gin.SetMode(gin.TestMode) + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequest("POST", "/api/v1/agents/chat/completions", + strings.NewReader(`{"agent_id":"a1","stream":true,"query":"hi"}`)) + c.Request.Header.Set("Content-Type", "application/json") + c.Set("user", &entity.User{ID: "u1"}) + c.Set("user_id", "u1") + + runner := &stubChatRunner{events: []canvas.RunEvent{ + {Type: "message", MessageID: "msg-1", TaskID: "task-1", SessionID: "sess-1", Data: `{"content":"hi back"}`}, + }} + h := &AgentHandler{chatRunner: runner} + h.AgentChatCompletions(c) + + body := w.Body.String() + if got := strings.Count(body, "data:[DONE]\n\n"); got != 1 { + t.Fatalf("expected exactly one [DONE] terminator, got %d in %q", got, body) + } + if !strings.HasSuffix(body, "data:[DONE]\n\n") { + t.Errorf("body should end with [DONE] terminator, got %q", body) + } +} + +func TestRunAgent_StreamAddsDoneWhenRunnerCloses(t *testing.T) { + gin.SetMode(gin.TestMode) + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Params = gin.Params{{Key: "canvas_id", Value: "a1"}} + c.Request = httptest.NewRequest("POST", "/api/v1/agents/a1/run", + strings.NewReader(`{"user_input":"hi"}`)) + c.Request.Header.Set("Content-Type", "application/json") + c.Set("user", &entity.User{ID: "u1"}) + c.Set("user_id", "u1") + + runner := &stubChatRunner{events: []canvas.RunEvent{ + {Type: "message", MessageID: "msg-1", TaskID: "task-1", SessionID: "sess-1", Data: `{"content":"hi back"}`}, + }} + h := &AgentHandler{chatRunner: runner} + h.RunAgent(c) + + body := w.Body.String() + if got := strings.Count(body, "data:[DONE]\n\n"); got != 1 { + t.Fatalf("expected exactly one [DONE] terminator, got %d in %q", got, body) + } + if !strings.HasSuffix(body, "data:[DONE]\n\n") { t.Errorf("body should end with [DONE] terminator, got %q", body) } } @@ -804,7 +855,7 @@ func TestAgentChatCompletions_DefaultBranchNonStreaming(t *testing.T) { !strings.Contains(body, `"hello back"`) { t.Errorf("body should contain agent event with content in data, got %q", body) } - if strings.Contains(body, "data: [DONE]") { + if strings.Contains(body, "data:[DONE]") { t.Errorf("body should not contain [DONE] terminator in non-streaming mode, got %q", body) } } diff --git a/internal/service/agent.go b/internal/service/agent.go index 14604def2f..09f5767ac2 100644 --- a/internal/service/agent.go +++ b/internal/service/agent.go @@ -88,6 +88,169 @@ func (s *AgentService) RunAgentWithWebhook( return s.RunAgent(ctx, userID, canvasID, "", "", "") } +func emitAgentMessageEvents(emit func(string, string), answer, thinking string, reference any) { + for _, ev := range buildAgentMessageEvents(answer, thinking, reference) { + data, _ := json.Marshal(ev) + emit("message", string(data)) + } +} + +type agentMessageDeltaEmitter struct { + emit func(string, string) + thinkState *ThinkStreamState + inThinking bool + explicitReasoning bool + emitted bool +} + +func newAgentMessageDeltaEmitter(emit func(string, string)) *agentMessageDeltaEmitter { + return &agentMessageDeltaEmitter{ + emit: emit, + thinkState: &ThinkStreamState{}, + } +} + +func (e *agentMessageDeltaEmitter) emitEvent(ev canvas.MessageEvent) { + emitAgentMessageEvent(e.emit, ev) + e.emitted = true +} + +func (e *agentMessageDeltaEmitter) startThinking() { + if e.inThinking { + return + } + e.emitEvent(canvas.MessageEvent{StartToThink: true}) + e.inThinking = true +} + +func (e *agentMessageDeltaEmitter) endThinking() { + if !e.inThinking { + return + } + e.emitEvent(canvas.MessageEvent{EndToThink: true}) + e.inThinking = false +} + +func (e *agentMessageDeltaEmitter) emitThinkDeltas(deltas []ThinkDelta) { + for _, d := range deltas { + switch { + case d.Kind == ThinkDeltaMarker && d.Value == thinkOpen: + e.startThinking() + case d.Kind == ThinkDeltaMarker && d.Value == thinkClose: + e.endThinking() + case d.Kind == ThinkDeltaText && d.Value != "": + e.emitEvent(canvas.MessageEvent{Content: d.Value}) + } + } +} + +func (e *agentMessageDeltaEmitter) Emit(contentDelta, thinkingDelta string) { + if thinkingDelta != "" { + e.startThinking() + e.explicitReasoning = true + e.emitEvent(canvas.MessageEvent{Content: thinkingDelta}) + } + if contentDelta == "" { + return + } + if e.explicitReasoning { + e.endThinking() + e.explicitReasoning = false + } + e.emitThinkDeltas(NextThinkDelta(e.thinkState, contentDelta, 0)) +} + +func (e *agentMessageDeltaEmitter) Finalize() bool { + before := e.emitted + e.emitThinkDeltas(FlushRemaining(e.thinkState)) + if e.explicitReasoning || e.inThinking { + e.endThinking() + e.explicitReasoning = false + } + return e.emitted && !before +} + +func (e *agentMessageDeltaEmitter) Reset() { + e.thinkState = &ThinkStreamState{} + e.inThinking = false + e.explicitReasoning = false + e.emitted = false +} + +func makeAgentMessageDeltaEmitter(emit func(string, string)) func(string, string) { + return newAgentMessageDeltaEmitter(emit).Emit +} + +func makeAgentMessageDeltaEmitterWithFinalizer(emit func(string, string)) (func(string, string), func() bool, func()) { + emitter := newAgentMessageDeltaEmitter(emit) + return emitter.Emit, emitter.Finalize, emitter.Reset +} + +func emitAgentMessageEvent(emit func(string, string), ev canvas.MessageEvent) { + data, _ := json.Marshal(ev) + emit("message", string(data)) +} + +func buildAgentMessageEvents(answer, thinking string, reference any) []canvas.MessageEvent { + answer, thinking = splitInlineThink(answer, thinking) + if thinking == "" { + return []canvas.MessageEvent{{ + Content: answer, + Reference: reference, + }} + } + + events := []canvas.MessageEvent{{StartToThink: true}} + for _, chunk := range splitMessageContent(thinking) { + events = append(events, canvas.MessageEvent{Content: chunk}) + } + events = append(events, canvas.MessageEvent{EndToThink: true}) + for _, chunk := range splitMessageContent(answer) { + events = append(events, canvas.MessageEvent{Content: chunk}) + } + return events +} + +func splitInlineThink(answer, thinking string) (string, string) { + if thinking != "" { + return answer, thinking + } + const startTag = "" + const endTag = "" + start := strings.Index(answer, startTag) + if start < 0 { + return answer, thinking + } + afterStart := start + len(startTag) + endRel := strings.Index(answer[afterStart:], endTag) + if endRel < 0 { + return answer, thinking + } + end := afterStart + endRel + thinking = answer[afterStart:end] + answer = answer[:start] + answer[end+len(endTag):] + answer = strings.TrimLeft(answer, "\r\n") + return answer, thinking +} + +func splitMessageContent(content string) []string { + if content == "" { + return nil + } + const maxRunes = 24 + runes := []rune(content) + chunks := make([]string, 0, (len(runes)+maxRunes-1)/maxRunes) + for len(runes) > 0 { + n := maxRunes + if len(runes) < n { + n = len(runes) + } + chunks = append(chunks, string(runes[:n])) + runes = runes[n:] + } + return chunks +} + // ErrAgentNotOwner is returned by DeleteAgent when the canvas exists and // is accessible to the caller but is owned by a different user. It maps // to the Python "Only the owner of the agent is authorized for this @@ -998,6 +1161,8 @@ func (s *AgentService) buildRunFunc(canvasID string, versionRow *entity.UserCanv TaskID: taskID, SessionID: sessionID, }) + agentMessageEmit, agentMessageFinalize, agentMessageReset := makeAgentMessageDeltaEmitterWithFinalizer(emit) + ctx2 = runtime.WithAgentMessageEmitterControl(ctx2, agentMessageEmit, agentMessageFinalize, agentMessageReset) // Seed initial env/sys values from the Canvas DSL globals. // Python's self.globals dict stores "sys.*" and "env.*" under @@ -1116,6 +1281,7 @@ func (s *AgentService) buildRunFunc(canvasID string, versionRow *entity.UserCanv // node_finished events are already emitted per-node by the // statePost wrappers in scheduler.go. var answer string + var thinking string var legacyReference []interface{} var downloads any now := float64(time.Now().UnixNano()) / 1e9 @@ -1131,6 +1297,9 @@ func (s *AgentService) buildRunFunc(canvasID string, versionRow *entity.UserCanv if v, ok := bucket["result"].(string); ok && v != "" && answer == "" { answer = v } + if v, ok := bucket["thinking"].(string); ok && v != "" && thinking == "" { + thinking = v + } if v, ok := bucket["reference"].([]interface{}); ok { legacyReference = append(legacyReference, v...) } @@ -1139,6 +1308,8 @@ func (s *AgentService) buildRunFunc(canvasID string, versionRow *entity.UserCanv } } referencePayload := agentRunReferencePayload(state, legacyReference) + runtime.FinalizeAgentMessage(ctx2) + messageEventsEmitted := runtime.AgentMessageEventsEmitted(ctx2) if err != nil { common.Debug("RunAgent invoke err", @@ -1152,11 +1323,9 @@ func (s *AgentService) buildRunFunc(canvasID string, versionRow *entity.UserCanv s.markRunFailed(ctx2, runID, "interrupt: "+err.Error()) if answer != "" { s.persistAgentRunSession(canvasID, sessionID, messageID, userInput, answer, referencePayload) - msgData, _ := json.Marshal(canvas.MessageEvent{ - Content: answer, - Reference: referencePayload, - }) - emit("message", string(msgData)) + if !messageEventsEmitted { + emitAgentMessageEvents(emit, answer, thinking, referencePayload) + } meData, _ := json.Marshal(canvas.MessageEndEvent{ Reference: referencePayload, @@ -1167,11 +1336,9 @@ func (s *AgentService) buildRunFunc(canvasID string, versionRow *entity.UserCanv } if shouldTreatAsCompletedLoopRun(err, answer) { s.persistAgentRunSession(canvasID, sessionID, messageID, userInput, answer, referencePayload) - msgData, _ := json.Marshal(canvas.MessageEvent{ - Content: answer, - Reference: referencePayload, - }) - emit("message", string(msgData)) + if !messageEventsEmitted { + emitAgentMessageEvents(emit, answer, thinking, referencePayload) + } meData, _ := json.Marshal(canvas.MessageEndEvent{ Reference: referencePayload, @@ -1199,11 +1366,9 @@ func (s *AgentService) buildRunFunc(canvasID string, versionRow *entity.UserCanv // Emit message + message_end (mirrors Python's ans dict). s.persistAgentRunSession(canvasID, sessionID, messageID, userInput, answer, referencePayload) - msgData, _ := json.Marshal(canvas.MessageEvent{ - Content: answer, - Reference: referencePayload, - }) - emit("message", string(msgData)) + if !messageEventsEmitted { + emitAgentMessageEvents(emit, answer, thinking, referencePayload) + } meData, _ := json.Marshal(canvas.MessageEndEvent{ Reference: referencePayload, diff --git a/internal/service/agent_run_e2e_test.go b/internal/service/agent_run_e2e_test.go index e4c9cd1ef4..4b0645807f 100644 --- a/internal/service/agent_run_e2e_test.go +++ b/internal/service/agent_run_e2e_test.go @@ -92,6 +92,7 @@ func drainAgentEvents(t *testing.T, events <-chan canvas.RunEvent) (messages []c select { case ev, ok := <-events: if !ok { + done = true return } switch ev.Type { @@ -434,8 +435,8 @@ func TestRunAgent_RealCanvas_WaitForUserResume_EventSemantics(t *testing.T) { t.Fatalf("run 1: unexpected workflow_finished before wait-for-user, events=%v", types1) } } - if len(types1) == 0 || types1[len(types1)-2] != "waiting_for_user" || types1[len(types1)-1] != "done" { - t.Fatalf("run 1: tail events = %v, want ... waiting_for_user, done", types1) + if len(types1) == 0 || types1[len(types1)-1] != "waiting_for_user" { + t.Fatalf("run 1: tail events = %v, want ... waiting_for_user", types1) } events2, err := svc.RunAgent( @@ -455,8 +456,8 @@ func TestRunAgent_RealCanvas_WaitForUserResume_EventSemantics(t *testing.T) { t.Fatalf("run 2: unexpected workflow_started on resume, events=%v", types2) } } - if len(types2) == 0 || types2[len(types2)-2] != "workflow_finished" || types2[len(types2)-1] != "done" { - t.Fatalf("run 2: tail events = %v, want ... workflow_finished, done", types2) + if len(types2) == 0 || types2[len(types2)-1] != "workflow_finished" { + t.Fatalf("run 2: tail events = %v, want ... workflow_finished", types2) } } diff --git a/internal/service/agent_test.go b/internal/service/agent_test.go index 3f93eb0ba5..aff57333d9 100644 --- a/internal/service/agent_test.go +++ b/internal/service/agent_test.go @@ -32,6 +32,221 @@ import ( "ragflow/internal/entity" ) +func TestBuildAgentMessageEventsThinkingProtocol(t *testing.T) { + events := buildAgentMessageEvents("final answer", "think step", nil) + if len(events) < 4 { + t.Fatalf("events len = %d, want at least start, thinking, end, answer", len(events)) + } + if !events[0].StartToThink { + t.Fatalf("first event StartToThink = false") + } + if events[0].Content != "" { + t.Fatalf("start event content = %q, want empty", events[0].Content) + } + + endIdx := -1 + for i, ev := range events { + if ev.EndToThink { + endIdx = i + break + } + } + if endIdx < 0 { + t.Fatal("missing EndToThink event") + } + var thinking strings.Builder + for _, ev := range events[1:endIdx] { + thinking.WriteString(ev.Content) + } + if got := thinking.String(); got != "think step" { + t.Fatalf("thinking stream = %q, want %q", got, "think step") + } + var answer strings.Builder + for _, ev := range events[endIdx+1:] { + answer.WriteString(ev.Content) + } + if got := answer.String(); got != "final answer" { + t.Fatalf("answer stream = %q, want %q", got, "final answer") + } +} + +func TestBuildAgentMessageEventsSplitsInlineThink(t *testing.T) { + events := buildAgentMessageEvents("plan\nanswer", "", nil) + if len(events) < 4 || !events[0].StartToThink { + t.Fatalf("inline think events malformed: %+v", events) + } + + endIdx := -1 + for i, ev := range events { + if ev.EndToThink { + endIdx = i + break + } + } + if endIdx < 0 { + t.Fatal("missing EndToThink event") + } + if got := events[1].Content; got != "plan" { + t.Fatalf("inline thinking = %q, want plan", got) + } + var answer strings.Builder + for _, ev := range events[endIdx+1:] { + answer.WriteString(ev.Content) + } + if got := answer.String(); got != "answer" { + t.Fatalf("inline answer = %q, want answer", got) + } +} + +func TestBuildAgentMessageEventsWithoutThinkingKeepsSingleMessage(t *testing.T) { + ref := map[string]any{"total": 1} + events := buildAgentMessageEvents("answer", "", ref) + if len(events) != 1 { + t.Fatalf("events len = %d, want 1", len(events)) + } + if events[0].Content != "answer" { + t.Fatalf("content = %q, want answer", events[0].Content) + } + if events[0].Reference == nil { + t.Fatal("reference missing") + } +} + +func TestAgentMessageDeltaEmitterStreamsInlineThink(t *testing.T) { + var events []canvas.MessageEvent + emit := makeAgentMessageDeltaEmitter(func(event, data string) { + if event != "message" { + t.Fatalf("event = %q, want message", event) + } + var ev canvas.MessageEvent + if err := json.Unmarshal([]byte(data), &ev); err != nil { + t.Fatalf("unmarshal event: %v", err) + } + events = append(events, ev) + }) + + emit("plan", "") + emit("answer", "") + + if len(events) != 4 { + t.Fatalf("events len = %d, want 4: %+v", len(events), events) + } + if !events[0].StartToThink { + t.Fatalf("first event = %+v, want StartToThink", events[0]) + } + if events[1].Content != "plan" { + t.Fatalf("thinking content = %q, want plan", events[1].Content) + } + if !events[2].EndToThink { + t.Fatalf("third event = %+v, want EndToThink", events[2]) + } + if events[3].Content != "answer" { + t.Fatalf("answer content = %q, want answer", events[3].Content) + } +} + +func TestAgentMessageDeltaEmitterStreamsReasoningBeforeAnswer(t *testing.T) { + var events []canvas.MessageEvent + emit := makeAgentMessageDeltaEmitter(func(event, data string) { + if event != "message" { + t.Fatalf("event = %q, want message", event) + } + var ev canvas.MessageEvent + if err := json.Unmarshal([]byte(data), &ev); err != nil { + t.Fatalf("unmarshal event: %v", err) + } + events = append(events, ev) + }) + + emit("", "step 1") + emit("", " step 2") + emit("answer", "") + + if len(events) != 5 { + t.Fatalf("events len = %d, want 5: %+v", len(events), events) + } + if !events[0].StartToThink { + t.Fatalf("first event = %+v, want StartToThink", events[0]) + } + if events[1].Content+events[2].Content != "step 1 step 2" { + t.Fatalf("thinking content = %q, want step 1 step 2", events[1].Content+events[2].Content) + } + if !events[3].EndToThink { + t.Fatalf("fourth event = %+v, want EndToThink", events[3]) + } + if events[4].Content != "answer" { + t.Fatalf("answer content = %q, want answer", events[4].Content) + } +} + +func TestAgentMessageDeltaEmitterProcessesThinkingAndContentTogether(t *testing.T) { + var events []canvas.MessageEvent + emit, finalize, _ := makeAgentMessageDeltaEmitterWithFinalizer(func(event, data string) { + if event != "message" { + t.Fatalf("event = %q, want message", event) + } + var ev canvas.MessageEvent + if err := json.Unmarshal([]byte(data), &ev); err != nil { + t.Fatalf("unmarshal event: %v", err) + } + events = append(events, ev) + }) + + emit("answer", "think") + finalize() + + if len(events) != 4 { + t.Fatalf("events len = %d, want 4: %+v", len(events), events) + } + if !events[0].StartToThink { + t.Fatalf("first event = %+v, want StartToThink", events[0]) + } + if events[1].Content != "think" { + t.Fatalf("thinking content = %q, want think", events[1].Content) + } + if !events[2].EndToThink { + t.Fatalf("third event = %+v, want EndToThink", events[2]) + } + if events[3].Content != "answer" { + t.Fatalf("answer content = %q, want answer", events[3].Content) + } +} + +func TestAgentMessageDeltaEmitterFinalizeClosesReasoning(t *testing.T) { + var events []canvas.MessageEvent + emit, finalize, _ := makeAgentMessageDeltaEmitterWithFinalizer(func(event, data string) { + if event != "message" { + t.Fatalf("event = %q, want message", event) + } + var ev canvas.MessageEvent + if err := json.Unmarshal([]byte(data), &ev); err != nil { + t.Fatalf("unmarshal event: %v", err) + } + events = append(events, ev) + }) + + emit("", "think only") + finalize() + + if len(events) != 3 { + t.Fatalf("events len = %d, want 3: %+v", len(events), events) + } + if !events[0].StartToThink { + t.Fatalf("first event = %+v, want StartToThink", events[0]) + } + if events[1].Content != "think only" { + t.Fatalf("thinking content = %q, want think only", events[1].Content) + } + if !events[2].EndToThink { + t.Fatalf("third event = %+v, want EndToThink", events[2]) + } +} + // TestListVersions_Success verifies that ListVersions returns all versions // for a canvas, ordered by update_time DESC. func TestListVersions_Success(t *testing.T) { @@ -439,11 +654,10 @@ func TestRunAgent_NoVersionPublishedPlaceholder(t *testing.T) { // answer text is present. The driver emits at least one // orchestrator (canvas.Runner) RunEvent with Type=="message" whose Data is a // JSON-encoded MessageEvent with the placeholder Content, plus - // a terminator RunEvent with Type=="done". + // the handler writes the final data:[DONE] frame after the channel closes. var ( gotAnswer string gotMessageEvent bool - gotDoneEvent bool ) deadline := time.After(5 * time.Second) for { @@ -462,9 +676,6 @@ func TestRunAgent_NoVersionPublishedPlaceholder(t *testing.T) { if !strings.Contains(gotAnswer, "No published version") { t.Errorf("placeholder answer %q does not contain 'No published version'", gotAnswer) } - if !gotDoneEvent { - t.Error("placeholder channel closed without emitting a DoneEvent") - } return } switch ev.Type { @@ -475,8 +686,6 @@ func TestRunAgent_NoVersionPublishedPlaceholder(t *testing.T) { t.Fatalf("message RunEvent had un-decodable Data %q: %v", ev.Data, err) } gotAnswer = msg.Content - case "done": - gotDoneEvent = true } case <-deadline: t.Fatal("placeholder channel did not close within 5s — driver deadlocked?") diff --git a/internal/service/bot_completion.go b/internal/service/bot_completion.go index da33565f42..65ab8bd7c5 100644 --- a/internal/service/bot_completion.go +++ b/internal/service/bot_completion.go @@ -169,7 +169,7 @@ func WriteDoneFrame(w http.ResponseWriter) error { // data.answer, the browser receives bytes but cannot render the // assistant message or correlate the current Log panel. // -// The "done" event type emits `data: [DONE]\n\n` (no envelope), +// The "done" event type emits `data:[DONE]\n\n` (no envelope), // matching the Python agent API terminator. // // Returns the write error so callers can short-circuit; both nil @@ -177,7 +177,7 @@ func WriteDoneFrame(w http.ResponseWriter) error { // disconnected mid-stream. func WriteChatbotRunEvent(w http.ResponseWriter, ev canvas.RunEvent) error { if ev.Type == "done" { - _, err := w.Write([]byte("data: [DONE]\n\n")) + _, err := w.Write([]byte("data:[DONE]\n\n")) if err != nil { return err }