diff --git a/internal/agent/canvas/loop_subgraph.go b/internal/agent/canvas/loop_subgraph.go index 30227877af..bf420eb9b8 100644 --- a/internal/agent/canvas/loop_subgraph.go +++ b/internal/agent/canvas/loop_subgraph.go @@ -40,6 +40,7 @@ import ( "slices" "strings" + "ragflow/internal/agent/runtime" "ragflow/internal/agent/workflowx" "github.com/cloudwego/eino/compose" @@ -241,7 +242,12 @@ func buildSubWorkflow( if name == "" { return nil, fmt.Errorf("canvas: loop %q member %q has empty component_name", loopID, cpnID) } - body, err := buildNodeBody(ctx, cpnID, name, c.Components[cpnID].Obj.Params) + deferToMessage := directMessageDownstream(c, cpnID) + nodeOpts := runtime.ComponentExecutionOptions{ + DeferAgentToMessage: deferToMessage, + SuppressAgentMessageEvents: strings.EqualFold(name, "Agent") && !deferToMessage, + } + body, err := buildNodeBodyWithOptions(ctx, cpnID, name, c.Components[cpnID].Obj.Params, nodeOpts) if err != nil { return nil, err } diff --git a/internal/agent/canvas/node_body.go b/internal/agent/canvas/node_body.go index 2eb93be504..8e78ec3135 100644 --- a/internal/agent/canvas/node_body.go +++ b/internal/agent/canvas/node_body.go @@ -120,6 +120,10 @@ func applyOverrideParams(params, cpnOverride map[string]any) map[string]any { } func buildNodeBody(ctx context.Context, cpnID, name string, params map[string]any) (nodeBodyFn, error) { + return buildNodeBodyWithOptions(ctx, cpnID, name, params, runtime.ComponentExecutionOptions{}) +} + +func buildNodeBodyWithOptions(ctx context.Context, cpnID, name string, params map[string]any, opts runtime.ComponentExecutionOptions) (nodeBodyFn, error) { if overrides := overrideParamsFromContext(ctx); len(overrides) > 0 { // overrides is keyed by cpnID; a component only sees its own // entry. Components absent from the map are left untouched. @@ -155,7 +159,7 @@ func buildNodeBody(ctx context.Context, cpnID, name string, params map[string]an // to expose Name(). The factory returns the class name as // the DSL's `component_name` field, which is also what // ComponentBase.Name() would have returned. - return realComponentBody(cpnID, name, comp), nil + return realComponentBodyWithOptions(cpnID, name, comp, opts), nil } // Fallback: no factory registered. This path is only exercised by // canvas-only unit tests; production wiring always installs a @@ -238,6 +242,10 @@ func componentTimeout() time.Duration { // key it is overwritten with the canvas-controlled value to keep // attribution authoritative. func realComponentBody(cpnID, componentClass string, comp runtime.Component) nodeBodyFn { + return realComponentBodyWithOptions(cpnID, componentClass, comp, runtime.ComponentExecutionOptions{}) +} + +func realComponentBodyWithOptions(cpnID, componentClass string, comp runtime.Component, opts runtime.ComponentExecutionOptions) nodeBodyFn { return func(ctx context.Context, in map[string]any) (map[string]any, error) { timeout := resolveTimeoutFromContext(ctx, componentClass) cctx, cancel := context.WithTimeout(ctx, timeout) @@ -247,7 +255,7 @@ func realComponentBody(cpnID, componentClass string, comp runtime.Component) nod invokeErr := runtime.TrackProgress(cpnID, runtime.ProgressCallbackFromContext(ctx), func() error { var e error out, e = runtime.TrackElapsed(componentClass, func() (map[string]any, error) { - return comp.Invoke(cctx, in) + return comp.Invoke(runtime.WithComponentExecutionOptions(cctx, opts), in) }) return e }) @@ -342,7 +350,13 @@ func withStateBracket(cpnID, componentName string, body nodeBodyFn) nodeBodyFn { } state.SetVar(outputCpnID, k, v) } - nodeFinishedNow(ctx, state, cpnID, componentName, componentName, nil) + if runtime.IsDeferredStream(out["content"]) { + runtime.RegisterDeferredNode(ctx, cpnID, func() { + nodeFinishedNow(ctx, state, cpnID, componentName, componentName, nil) + }) + } else { + nodeFinishedNow(ctx, state, cpnID, componentName, componentName, nil) + } return out, nil } } diff --git a/internal/agent/canvas/scheduler.go b/internal/agent/canvas/scheduler.go index e5a2df82be..19f3b1d20a 100644 --- a/internal/agent/canvas/scheduler.go +++ b/internal/agent/canvas/scheduler.go @@ -549,7 +549,12 @@ func BuildWorkflow(ctx context.Context, c *Canvas) (*compose.Workflow[map[string if name == "" { return nil, fmt.Errorf("canvas: component %q has empty component_name", cpnID) } - body, err := buildNodeBody(ctx, cpnID, name, c.Components[cpnID].Obj.Params) + deferToMessage := directMessageDownstream(c, cpnID) + nodeOpts := runtime.ComponentExecutionOptions{ + DeferAgentToMessage: deferToMessage, + SuppressAgentMessageEvents: strings.EqualFold(name, "Agent") && !deferToMessage, + } + body, err := buildNodeBodyWithOptions(ctx, cpnID, name, c.Components[cpnID].Obj.Params, nodeOpts) if err != nil { return nil, err } @@ -566,7 +571,16 @@ func BuildWorkflow(ctx context.Context, c *Canvas) (*compose.Workflow[map[string } nodePost := func(ctx context.Context, out map[string]any, state *CanvasState) (map[string]any, error) { result, postErr := statePost(ctx, out, state) - nodeFinishedNow(ctx, state, cpnID, componentName, componentName, postErr) + if postErr == nil && runtime.IsDeferredStream(result["content"]) { + // Python keeps the Agent node pending while Message consumes its + // partial generator. Message completes this callback after the + // deferred stream closes. + runtime.RegisterDeferredNode(ctx, cpnID, func() { + nodeFinishedNow(ctx, state, cpnID, componentName, componentName, nil) + }) + } else { + nodeFinishedNow(ctx, state, cpnID, componentName, componentName, postErr) + } return result, postErr } lambda := compose.InvokableLambda[map[string]any, map[string]any](body) @@ -699,6 +713,26 @@ func BuildWorkflow(ctx context.Context, c *Canvas) (*compose.Workflow[map[string return wf, nil } +// directMessageDownstream: only a direct +// Message child enables lazy Agent execution. Intermediate nodes must not +// accidentally change the Agent's execution mode. +func directMessageDownstream(c *Canvas, cpnID string) bool { + if c == nil { + return false + } + comp, ok := c.Components[cpnID] + if !ok { + return false + } + for _, downID := range comp.Downstream { + down, ok := c.Components[downID] + if ok && strings.EqualFold(down.Obj.ComponentName, "Message") { + return true + } + } + return false +} + func wireWorkflowTerminals( wf *compose.Workflow[map[string]any, map[string]any], terminals []string, diff --git a/internal/agent/component/agent.go b/internal/agent/component/agent.go index 816b0ac28a..e9844d9c12 100644 --- a/internal/agent/component/agent.go +++ b/internal/agent/component/agent.go @@ -19,6 +19,7 @@ import ( "io" "sort" "strings" + "time" einotool "github.com/cloudwego/eino/components/tool" "github.com/cloudwego/eino/compose" @@ -35,6 +36,7 @@ import ( ) const maxSubAgentDepth = 8 +const defaultAgentDeferredTimeout = 10 * time.Minute // agentPatternless matches `@` and // `@@` (the trailing `@` is @@ -282,7 +284,7 @@ func emitAgentModelStreams(ctx context.Context, future react.MessageFuture) <-ch if msg.Content == "" && msg.ReasoningContent == "" { continue } - if runtime.AgentMessageEventsEmitted(ctx) { + if runtime.AgentMessageEventsEmitted(ctx) && !runtime.HasDeferredAgentMessageSink(ctx) { continue } runtime.EmitAgentMessage(ctx, msg.Content, msg.ReasoningContent) @@ -710,9 +712,35 @@ func NewAgentComponent(p AgentParam) *AgentComponent { // Name returns the registered component name. func (c *AgentComponent) Name() string { return "Agent" } -// Invoke runs the ReAct loop via the configured agentRunner and returns -// the output map. +// Invoke either returns a lazy Agent stream for a direct downstream Message, +// or executes the Agent eagerly for all other graph shapes. The mode is a +// compile-time canvas decision carried through context, not a DSL parameter. func (c *AgentComponent) Invoke(ctx context.Context, inputs map[string]any) (map[string]any, error) { + if runtime.ComponentExecutionOptionsFromContext(ctx).DeferAgentToMessage { + // Preserve the Agent-class duration installed by the node wrapper, then + // start a fresh deadline when Message opens the lazy stream. + timeout := defaultAgentDeferredTimeout + if deadline, ok := ctx.Deadline(); ok { + if remaining := time.Until(deadline); remaining > 0 { + timeout = remaining + } + } + deferred := &runtime.DeferredStream{ + Open: func(openCtx context.Context, sink runtime.AgentDeltaSink) (map[string]any, error) { + agentCtx, cancel := context.WithTimeout(openCtx, timeout) + defer cancel() + return c.invokeNow(runtime.WithAgentDeltaSink(agentCtx, sink), inputs) + }, + } + return map[string]any{"content": deferred}, nil + } + return c.invokeNow(ctx, inputs) +} + +// invokeNow contains the original eager Agent execution path. Deferred +// Message consumption calls this function later with an invocation-local +// delta sink, so the same ReAct/citation/tool behavior is reused once. +func (c *AgentComponent) invokeNow(ctx context.Context, inputs map[string]any) (map[string]any, error) { runtime.ResetAgentMessageEmission(ctx) defer runtime.FinalizeAgentMessage(ctx) @@ -852,7 +880,8 @@ func (c *AgentComponent) Invoke(ctx context.Context, inputs map[string]any) (map if groundingStatus != "" { out["grounding_status"] = groundingStatus } - if !runtime.AgentMessageEventsEmitted(ctx) { + streamed := runtime.AgentMessageEventsEmitted(ctx) || runtime.DeferredAgentMessageEventsEmitted(ctx) + if !streamed { runtime.EmitAgentMessage(ctx, content+artifactMD, thinking) } return out, nil @@ -961,7 +990,7 @@ func buildAgentChatModel(ctx context.Context, p AgentParam) (*models.EinoChatMod // would be dead weight. When AgentParam grows Temperature/ // MaxTokens, switch to always-build. var chatCfg *models.ChatConfig - if p.TopP != nil || p.Thinking != "" || runtime.HasAgentMessageEmitter(ctx) { + if p.TopP != nil || p.Thinking != "" { chatCfg = &models.ChatConfig{TopP: p.TopP} switch p.Thinking { case "enabled": @@ -971,11 +1000,6 @@ func buildAgentChatModel(ctx context.Context, p AgentParam) (*models.EinoChatMod 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 } diff --git a/internal/agent/component/agent_test.go b/internal/agent/component/agent_test.go index b51b395639..e7c67ccd62 100644 --- a/internal/agent/component/agent_test.go +++ b/internal/agent/component/agent_test.go @@ -132,6 +132,131 @@ func TestAgent_MessageEmissionIsScopedPerInvocation(t *testing.T) { } } +func TestAgent_DefersExecutionForDownstreamMessage(t *testing.T) { + calls := 0 + withAgentRunner(t, func(_ context.Context, _ AgentParam) (*schema.Message, error) { + calls++ + return &schema.Message{Role: schema.Assistant, Content: "lazy answer"}, nil + }) + + ctx := runtime.WithComponentExecutionOptions(context.Background(), runtime.ComponentExecutionOptions{ + DeferAgentToMessage: true, + }) + ctx = runtime.WithAgentMessageEmitter(ctx, func(string, string) {}) + agent := NewAgentComponent(AgentParam{ModelID: "stub", MaxRounds: 1}) + out, err := agent.Invoke(ctx, map[string]any{"user_prompt": "hello"}) + if err != nil { + t.Fatalf("Invoke: %v", err) + } + if calls != 0 { + t.Fatalf("runner calls=%d before Message consumed stream, want 0", calls) + } + deferred, ok := out["content"].(*runtime.DeferredStream) + if !ok || deferred == nil { + t.Fatalf("content=%T, want *runtime.DeferredStream", out["content"]) + } + var got strings.Builder + final, err := deferred.Open(ctx, func(contentDelta, _ string) { + got.WriteString(contentDelta) + }) + if err != nil { + t.Fatalf("Open: %v", err) + } + if calls != 1 || got.String() != "lazy answer" { + t.Fatalf("calls=%d streamed=%q, want 1 / %q", calls, got.String(), "lazy answer") + } + if final["content"] != "lazy answer" { + t.Fatalf("final content=%v, want lazy answer", final["content"]) + } +} + +func TestAgent_DeferredStreamAppliesFreshTimeout(t *testing.T) { + withAgentRunner(t, func(ctx context.Context, _ AgentParam) (*schema.Message, error) { + if _, ok := ctx.Deadline(); !ok { + t.Fatal("deferred Agent runner context has no deadline") + } + return &schema.Message{Role: schema.Assistant, Content: "answer"}, nil + }) + + ctx := runtime.WithComponentExecutionOptions(context.Background(), runtime.ComponentExecutionOptions{ + DeferAgentToMessage: true, + }) + out, err := NewAgentComponent(AgentParam{ModelID: "stub"}).Invoke(ctx, map[string]any{"user_prompt": "hello"}) + if err != nil { + t.Fatalf("Invoke: %v", err) + } + deferred, ok := out["content"].(*runtime.DeferredStream) + if !ok || deferred == nil { + t.Fatalf("content=%T, want *runtime.DeferredStream", out["content"]) + } + if _, err := deferred.Open(context.Background(), func(string, string) {}); err != nil { + t.Fatalf("Open: %v", err) + } +} + +func TestAgent_DeferredStreamDoesNotAppendFinalAnswerAfterDeltas(t *testing.T) { + withAgentRunner(t, func(ctx context.Context, _ AgentParam) (*schema.Message, error) { + runtime.EmitAgentMessage(ctx, "lazy ", "") + runtime.EmitAgentMessage(ctx, "answer", "") + return &schema.Message{Role: schema.Assistant, Content: "lazy answer"}, nil + }) + + ctx := runtime.WithComponentExecutionOptions(context.Background(), runtime.ComponentExecutionOptions{ + DeferAgentToMessage: true, + }) + ctx = runtime.WithAgentMessageEmitter(ctx, func(string, string) {}) + out, err := NewAgentComponent(AgentParam{ModelID: "stub", MaxRounds: 1}).Invoke(ctx, map[string]any{ + "user_prompt": "hello", + }) + if err != nil { + t.Fatalf("Invoke: %v", err) + } + deferred, ok := out["content"].(*runtime.DeferredStream) + if !ok || deferred == nil { + t.Fatalf("content=%T, want *runtime.DeferredStream", out["content"]) + } + + var streamed strings.Builder + final, err := deferred.Open(ctx, func(contentDelta, _ string) { + streamed.WriteString(contentDelta) + }) + if err != nil { + t.Fatalf("Open: %v", err) + } + if got, want := streamed.String(), "lazy answer"; got != want { + t.Fatalf("streamed=%q, want %q", got, want) + } + if got, want := final["content"], "lazy answer"; got != want { + t.Fatalf("final content=%v, want %v", got, want) + } +} + +func TestAgent_SuppressesVisibleEventsWithoutMessageDownstream(t *testing.T) { + withAgentRunner(t, func(_ context.Context, _ AgentParam) (*schema.Message, error) { + return &schema.Message{Role: schema.Assistant, Content: "timeline answer"}, nil + }) + var emitted []string + ctx := runtime.WithAgentMessageEmitter(context.Background(), func(contentDelta, _ string) { + if contentDelta != "" { + emitted = append(emitted, contentDelta) + } + }) + ctx = runtime.WithComponentExecutionOptions(ctx, runtime.ComponentExecutionOptions{ + SuppressAgentMessageEvents: true, + }) + if _, err := NewAgentComponent(AgentParam{ModelID: "stub", MaxRounds: 1}).Invoke(ctx, map[string]any{ + "user_prompt": "hello", + }); err != nil { + t.Fatalf("Invoke: %v", err) + } + if len(emitted) != 0 { + t.Fatalf("visible events=%#v, want none", emitted) + } + if !runtime.AgentMessageEventsSuppressed(ctx) { + t.Fatal("AgentMessageEventsSuppressed=false, want true") + } +} + func TestAgent_ForwardsThinkingParam(t *testing.T) { var gotThinking string withAgentRunner(t, func(_ context.Context, p AgentParam) (*schema.Message, error) { diff --git a/internal/agent/component/message.go b/internal/agent/component/message.go index d17b33cf02..58e4e542f6 100644 --- a/internal/agent/component/message.go +++ b/internal/agent/component/message.go @@ -186,10 +186,14 @@ func (m *MessageComponent) Invoke(ctx context.Context, inputs map[string]any) (m text = fallbackMessageText(inputs) } - // Message renders values for display, so keep Python's tolerant behavior: - // missing references become empty strings and structured values use JSON. - // Parameter-binding components continue to use the strict resolver. - resolved := runtime.ResolveTemplateForDisplay(text, state) + // A direct Agent→Message edge stores a lazy DeferredStream in the Agent + // output. Message is the owner of that stream: opening it here preserves + // Python's partial(async_generator) execution order and makes this node the + // only visible SSE producer. + resolved, streamed, streamErr := m.resolveDeferredTemplate(ctx, text, state) + if streamErr != nil { + return nil, streamErr + } // Extract downloads. Walks inputs for download-info maps so // callers can attach binaries to the message body. @@ -222,16 +226,11 @@ func (m *MessageComponent) Invoke(ctx context.Context, inputs map[string]any) (m Text: resolved, }) } - // Skip emission when an upstream Agent component already streamed its - // answer via the agent message emitter. Without this guard the user sees - // the same content twice: once from the Agent's StreamCallback during the - // ReAct loop, and again from the Message node's own emit call here. - // The rendered content is still stored in outputs["content"] for state - // persistence and the service-layer answer collector. - if rendered != "" && !runtime.AgentMessageEventsEmitted(ctx) { - if !runtime.EmitCanvasMessage(ctx, rendered) { - runtime.EmitAgentMessage(ctx, rendered, "") - } + // The runtime emitter owns Agent-to-Message de-duplication. It suppresses + // only an exact copy of content already streamed by an upstream Agent, so a + // Message node that intentionally transforms the answer is still visible. + if rendered != "" && !streamed { + runtime.EmitCanvasMessage(ctx, rendered) } // Python's Message output schema always contains downloads, including an @@ -329,6 +328,101 @@ func (m *MessageComponent) Invoke(ctx context.Context, inputs map[string]any) (m return out, nil } +// resolveDeferredTemplate resolves a Message template while consuming any +// lazy Agent stream it references. It returns the complete visible text and a +// flag indicating whether a DeferredStream was opened. +func (m *MessageComponent) resolveDeferredTemplate(ctx context.Context, text string, state *runtime.CanvasState) (string, bool, error) { + matches := runtime.VarRefPattern.FindAllStringSubmatchIndex(text, -1) + if len(matches) == 0 { + return text, false, nil + } + // Ordinary Message templates are rendered and emitted once by Invoke. + // Only templates that actually reference a DeferredStream belong to the + // incremental presentation path below. Emitting literals/normal variable + // values here and then emitting the fully rendered string in Invoke would + // produce duplicate SSE message events for every non-deferred template. + hasDeferred := false + for _, match := range matches { + ref := text[match[2]:match[3]] + value, _ := state.GetVar(ref) + if runtime.IsDeferredStream(value) { + hasDeferred = true + break + } + } + if !hasDeferred { + return runtime.ResolveTemplateForDisplay(text, state), false, nil + } + var out strings.Builder + last := 0 + streamed := false + for _, match := range matches { + start, end := match[0], match[1] + refStart, refEnd := match[2], match[3] + literal := text[last:start] + if literal != "" { + runtime.EmitCanvasMessageEvent(ctx, literal, false, false) + out.WriteString(literal) + } + ref := text[refStart:refEnd] + value, _ := state.GetVar(ref) + deferred, ok := value.(*runtime.DeferredStream) + if !ok || deferred == nil || deferred.Open == nil { + resolved := runtime.ResolveTemplateForDisplay(text[start:end], state) + runtime.EmitCanvasMessageEvent(ctx, resolved, false, false) + out.WriteString(resolved) + last = end + continue + } + + streamed = true + inThinking := false + visible := strings.Builder{} + result, err := deferred.Open(ctx, func(contentDelta, reasoningDelta string) { + if reasoningDelta != "" { + if !inThinking { + runtime.EmitCanvasMessageEvent(ctx, "", true, false) + inThinking = true + } + runtime.EmitCanvasMessageEvent(ctx, reasoningDelta, false, false) + } + if contentDelta != "" { + if inThinking { + runtime.EmitCanvasMessageEvent(ctx, "", false, true) + inThinking = false + } + runtime.EmitCanvasMessageEvent(ctx, contentDelta, false, false) + visible.WriteString(contentDelta) + } + }) + if inThinking { + runtime.EmitCanvasMessageEvent(ctx, "", false, true) + } + if err != nil { + return "", true, fmt.Errorf("Message: consume deferred Agent stream: %w", err) + } + finalText := visible.String() + if result != nil { + if completedContent, ok := result["content"].(string); ok { + finalText = completedContent + } + } + if strings.Contains(ref, "@") { + parts := strings.SplitN(ref, "@", 2) + state.SetVar(parts[0], parts[1], finalText) + runtime.CompleteDeferredNode(ctx, parts[0]) + } + out.WriteString(finalText) + last = end + } + if last < len(text) { + tail := text[last:] + runtime.EmitCanvasMessageEvent(ctx, tail, false, false) + out.WriteString(tail) + } + return out.String(), streamed, nil +} + // extractMemoryIDs normalises a memory_ids value from inputs / // params. Accepts []string and []any[string]. func extractMemoryIDs(inputs map[string]any) []string { diff --git a/internal/agent/component/message_test.go b/internal/agent/component/message_test.go index 5d686cd57c..0fe16d26e3 100644 --- a/internal/agent/component/message_test.go +++ b/internal/agent/component/message_test.go @@ -18,6 +18,7 @@ package component import ( "context" + "strings" "testing" "ragflow/internal/agent/canvas" @@ -193,11 +194,33 @@ func TestMessage_EmitsDirectCanvasMessage(t *testing.T) { } } +func TestMessage_NormalTemplateEmitsOnlyRenderedMessage(t *testing.T) { + c, _ := NewMessageComponent(nil) + state := canvas.NewCanvasState("run-normal-template", "task-normal-template") + state.Sys["query"] = "world" + ctx := withStateForTest(context.Background(), state) + var direct []string + ctx = runtime.WithCanvasMessageEmitter(ctx, func(content string) { + direct = append(direct, content) + }) + + out, err := c.Invoke(ctx, map[string]any{"text": "hello {{sys.query}}"}) + if err != nil { + t.Fatalf("Invoke: %v", err) + } + if got, _ := out["content"].(string); got != "hello world" { + t.Fatalf("content: got %q, want %q", got, "hello world") + } + if len(direct) != 1 || direct[0] != "hello world" { + t.Fatalf("direct = %#v, want one rendered message", direct) + } +} + // TestMessage_SkipsEmissionWhenAgentAlreadyStreamed verifies that the // Message component does not re-emit content when an upstream Agent // component already streamed its answer. This prevents the double-reply -// bug (Agent → Message): the Agent streams via StreamCallback during the -// ReAct loop, and the Message node must not send the same content again. +// bug (Agent → Message): the Agent's model stream reaches the runtime before +// the Message node, and the Message node must not send the same content again. func TestMessage_SkipsEmissionWhenAgentAlreadyStreamed(t *testing.T) { c, _ := NewMessageComponent(nil) state := canvas.NewCanvasState("run-skip", "task-skip") @@ -208,9 +231,14 @@ func TestMessage_SkipsEmissionWhenAgentAlreadyStreamed(t *testing.T) { emitted = append(emitted, contentDelta) } }) + var direct []string + ctx = runtime.WithCanvasMessageEmitter(ctx, func(content string) { + direct = append(direct, content) + }) - // Simulate the Agent having already streamed its answer. - runtime.EmitAgentMessage(ctx, "agent streamed this", "") + // Simulate the Agent having already streamed its answer in deltas. + runtime.EmitAgentMessage(ctx, "agent ", "") + runtime.EmitAgentMessage(ctx, "streamed this", "") emitted = nil // clear setup emission so we only capture Message's output out, err := c.Invoke(ctx, map[string]any{"content": "agent streamed this"}) @@ -225,6 +253,133 @@ func TestMessage_SkipsEmissionWhenAgentAlreadyStreamed(t *testing.T) { if len(emitted) != 0 { t.Fatalf("emitted = %#v, want empty (Agent already streamed)", emitted) } + if len(direct) != 0 { + t.Fatalf("direct = %#v, want empty (Agent already streamed)", direct) + } +} + +func TestMessage_EmitsContentDifferentFromAgentStream(t *testing.T) { + c, _ := NewMessageComponent(nil) + state := canvas.NewCanvasState("run-distinct", "task-distinct") + ctx := withStateForTest(context.Background(), state) + var emitted []string + ctx = runtime.WithAgentMessageEmitter(ctx, func(contentDelta, thinkingDelta string) { + if contentDelta != "" { + emitted = append(emitted, contentDelta) + } + }) + var direct []string + ctx = runtime.WithCanvasMessageEmitter(ctx, func(content string) { + direct = append(direct, content) + }) + + runtime.EmitAgentMessage(ctx, "draft answer", "") + emitted = nil + + out, err := c.Invoke(ctx, map[string]any{"content": "final answer"}) + if err != nil { + t.Fatalf("Invoke: %v", err) + } + if got, _ := out["content"].(string); got != "final answer" { + t.Fatalf("content: got %q, want %q", got, "final answer") + } + if len(emitted) != 0 { + t.Fatalf("agent fallback emitted = %#v, want empty", emitted) + } + if len(direct) != 1 || direct[0] != "final answer" { + t.Fatalf("direct = %#v, want [final answer]", direct) + } +} + +func TestMessage_ConsumesDeferredAgentStream(t *testing.T) { + c, _ := NewMessageComponent(nil) + state := canvas.NewCanvasState("run-deferred", "task-deferred") + state.SetVar("agent_0", "content", &runtime.DeferredStream{ + Open: func(_ context.Context, sink runtime.AgentDeltaSink) (map[string]any, error) { + sink("hello ", "") + sink("world", "") + return map[string]any{"content": "hello world"}, nil + }, + }) + ctx := withStateForTest(context.Background(), state) + var emitted []string + ctx = runtime.WithCanvasMessageEmitter(ctx, func(content string) { + if content != "" { + emitted = append(emitted, content) + } + }) + + out, err := c.Invoke(ctx, map[string]any{"text": "answer: {{agent_0@content}}"}) + if err != nil { + t.Fatalf("Invoke: %v", err) + } + if got, _ := out["content"].(string); got != "answer: hello world" { + t.Fatalf("content: got %q, want %q", got, "answer: hello world") + } + if got := strings.Join(emitted, ""); got != "answer: hello world" { + t.Fatalf("emitted: got %q, want %q (%#v)", got, "answer: hello world", emitted) + } +} + +func TestMessage_DeferredStreamThinkingEvents(t *testing.T) { + c, _ := NewMessageComponent(nil) + state := canvas.NewCanvasState("run-deferred-thinking", "task-deferred-thinking") + state.SetVar("agent_0", "content", &runtime.DeferredStream{ + Open: func(_ context.Context, sink runtime.AgentDeltaSink) (map[string]any, error) { + sink("", "plan") + sink("answer", "") + return map[string]any{"content": "answer"}, nil + }, + }) + ctx := withStateForTest(context.Background(), state) + var events []string + ctx = runtime.WithCanvasMessageEventEmitter(ctx, func(content string, startToThink, endToThink bool) { + switch { + case startToThink: + events = append(events, "start") + case endToThink: + events = append(events, "end") + default: + events = append(events, content) + } + }) + + if _, err := c.Invoke(ctx, map[string]any{"text": "{{agent_0@content}}"}); err != nil { + t.Fatalf("Invoke: %v", err) + } + if got, want := strings.Join(events, "|"), "start|plan|end|answer"; got != want { + t.Fatalf("events=%q, want %q", got, want) + } +} + +func TestMessage_DeferredStreamUsesCompletedContent(t *testing.T) { + c, _ := NewMessageComponent(nil) + state := canvas.NewCanvasState("run-deferred-final", "task-deferred-final") + state.SetVar("agent_0", "content", &runtime.DeferredStream{ + Open: func(_ context.Context, sink runtime.AgentDeltaSink) (map[string]any, error) { + sink("raw answer", "") + return map[string]any{"content": "grounded answer [ID:1]"}, nil + }, + }) + ctx := withStateForTest(context.Background(), state) + var streamed []string + ctx = runtime.WithCanvasMessageEmitter(ctx, func(content string) { + streamed = append(streamed, content) + }) + + out, err := c.Invoke(ctx, map[string]any{"text": "{{agent_0@content}}"}) + if err != nil { + t.Fatalf("Invoke: %v", err) + } + if got, _ := out["content"].(string); got != "grounded answer [ID:1]" { + t.Fatalf("content: got %q, want completed Agent content", got) + } + if got, _ := state.GetVar("agent_0@content"); got != "grounded answer [ID:1]" { + t.Fatalf("state content: got %v, want completed Agent content", got) + } + if got := strings.Join(streamed, ""); got != "raw answer" { + t.Fatalf("streamed events: got %q, want live delta", got) + } } func TestMessage_FormalizedContentFallback(t *testing.T) { diff --git a/internal/agent/runtime/context.go b/internal/agent/runtime/context.go index d0c7f04e7d..05a44f08c3 100644 --- a/internal/agent/runtime/context.go +++ b/internal/agent/runtime/context.go @@ -28,6 +28,7 @@ package runtime import ( "context" "fmt" + "strings" "sync" ) @@ -44,12 +45,127 @@ type canvasMessageEmitterCtxKey struct{} // shape free of canvas/service imports. type AgentMessageEmitter func(contentDelta, thinkingDelta string) type CanvasMessageEmitter func(content string) +type CanvasMessageEventEmitter func(content string, startToThink, endToThink bool) +type AgentDeltaSink func(contentDelta, thinkingDelta string) + +// DeferredStream is a lazy component value. The producer is deliberately +// opened by the consuming Message node. +type DeferredStream struct { + Open func(context.Context, AgentDeltaSink) (map[string]any, error) +} + +func IsDeferredStream(v any) bool { + deferred, ok := v.(*DeferredStream) + return ok && deferred != nil && deferred.Open != nil +} type agentMessageEmitterState struct { - emit AgentMessageEmitter - finalize func() bool - reset func() - emitted bool + mu sync.Mutex + emit AgentMessageEmitter + finalize func() bool + reset func() + emitted bool + suppressed bool + runEmitted bool + runSuppressed bool + agentContent strings.Builder +} + +type agentDeltaSinkCtxKey struct{} +type canvasMessageEventEmitterCtxKey struct{} + +type agentDeltaSinkState struct { + mu sync.Mutex + sink AgentDeltaSink + emitted bool +} + +type componentExecutionOptionsCtxKey struct{} +type deferredNodeRegistryCtxKey struct{} + +type deferredNodeRegistry struct { + mu sync.Mutex + completions map[string]func() +} + +// ComponentExecutionOptions carries compile-time graph decisions into a +// component invocation without leaking internal flags into DSL parameters. +type ComponentExecutionOptions struct { + DeferAgentToMessage bool + SuppressAgentMessageEvents bool +} + +func WithComponentExecutionOptions(ctx context.Context, opts ComponentExecutionOptions) context.Context { + return context.WithValue(ctx, componentExecutionOptionsCtxKey{}, opts) +} + +func ComponentExecutionOptionsFromContext(ctx context.Context) ComponentExecutionOptions { + opts, _ := ctx.Value(componentExecutionOptionsCtxKey{}).(ComponentExecutionOptions) + return opts +} + +// WithDeferredNodeRegistry installs the per-run callbacks used to delay an +// Agent node_finished event until its downstream Message has consumed the +// lazy stream. +func WithDeferredNodeRegistry(ctx context.Context) context.Context { + return context.WithValue(ctx, deferredNodeRegistryCtxKey{}, &deferredNodeRegistry{ + completions: make(map[string]func()), + }) +} + +func RegisterDeferredNode(ctx context.Context, nodeID string, complete func()) { + registry, _ := ctx.Value(deferredNodeRegistryCtxKey{}).(*deferredNodeRegistry) + if registry == nil || nodeID == "" || complete == nil { + return + } + registry.mu.Lock() + registry.completions[nodeID] = complete + registry.mu.Unlock() +} + +func CompleteDeferredNode(ctx context.Context, nodeID string) { + registry, _ := ctx.Value(deferredNodeRegistryCtxKey{}).(*deferredNodeRegistry) + if registry == nil { + return + } + registry.mu.Lock() + complete := registry.completions[nodeID] + delete(registry.completions, nodeID) + registry.mu.Unlock() + if complete != nil { + complete() + } +} + +func CompleteAllDeferredNodes(ctx context.Context) { + registry, _ := ctx.Value(deferredNodeRegistryCtxKey{}).(*deferredNodeRegistry) + if registry == nil { + return + } + registry.mu.Lock() + callbacks := make([]func(), 0, len(registry.completions)) + for nodeID, complete := range registry.completions { + callbacks = append(callbacks, complete) + delete(registry.completions, nodeID) + } + registry.mu.Unlock() + for _, complete := range callbacks { + if complete != nil { + complete() + } + } +} + +func WithAgentDeltaSink(ctx context.Context, sink AgentDeltaSink) context.Context { + if sink == nil { + return ctx + } + return context.WithValue(ctx, agentDeltaSinkCtxKey{}, &agentDeltaSinkState{sink: sink}) +} + +func agentDeltaSinkFromContext(ctx context.Context) *agentDeltaSinkState { + sink, _ := ctx.Value(agentDeltaSinkCtxKey{}).(*agentDeltaSinkState) + return sink } // WithState attaches *CanvasState to ctx for retrieval by @@ -96,6 +212,15 @@ func WithCanvasMessageEmitter(ctx context.Context, emit CanvasMessageEmitter) co return context.WithValue(ctx, canvasMessageEmitterCtxKey{}, emit) } +// WithCanvasMessageEventEmitter installs the Message presentation callback +// used for content and thinking-boundary-aware events. +func WithCanvasMessageEventEmitter(ctx context.Context, emit CanvasMessageEventEmitter) context.Context { + if emit == nil { + return ctx + } + return context.WithValue(ctx, canvasMessageEventEmitterCtxKey{}, emit) +} + // HasAgentMessageEmitter reports whether the service layer installed an // Agent message stream callback on ctx. func HasAgentMessageEmitter(ctx context.Context) bool { @@ -106,31 +231,151 @@ func HasAgentMessageEmitter(ctx context.Context) bool { // 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 { + if sinkState := agentDeltaSinkFromContext(ctx); sinkState != nil { + sinkState.mu.Lock() + sink := sinkState.sink + sinkState.mu.Unlock() + if sink == nil { + return false + } + // A deferred Agent is being consumed by Message. Do not call the + // service SSE emitter here; Message owns the visible event stream. + sink(contentDelta, thinkingDelta) + sinkState.mu.Lock() + if contentDelta != "" || thinkingDelta != "" { + sinkState.emitted = true + } + sinkState.mu.Unlock() + return true + } state, ok := ctx.Value(agentMessageEmitterCtxKey{}).(*agentMessageEmitterState) - if !ok || state == nil || state.emit == nil { + if !ok || state == nil { return false } - state.emit(contentDelta, thinkingDelta) + if ComponentExecutionOptionsFromContext(ctx).SuppressAgentMessageEvents { + state.mu.Lock() + state.suppressed = true + state.runSuppressed = true + state.mu.Unlock() + return true + } + state.mu.Lock() + emit := state.emit + state.mu.Unlock() + if emit == nil { + return false + } + emit(contentDelta, thinkingDelta) + state.mu.Lock() if contentDelta != "" || thinkingDelta != "" { state.emitted = true + state.runEmitted = true } + if contentDelta != "" { + state.agentContent.WriteString(contentDelta) + } + state.mu.Unlock() return true } -// EmitCanvasMessage emits already-rendered Message-component content directly. -// It also marks the shared message stream as emitted so the service layer does -// not re-emit the final state content after workflow completion. +func AgentMessageEventsSuppressed(ctx context.Context) bool { + state, ok := ctx.Value(agentMessageEmitterCtxKey{}).(*agentMessageEmitterState) + if !ok || state == nil { + return false + } + state.mu.Lock() + defer state.mu.Unlock() + return state.suppressed +} + +// AgentMessageEventsEmittedRun reports whether any visible Agent or Message +// content was emitted during the entire canvas run. +func AgentMessageEventsEmittedRun(ctx context.Context) bool { + state, ok := ctx.Value(agentMessageEmitterCtxKey{}).(*agentMessageEmitterState) + if !ok || state == nil { + return false + } + state.mu.Lock() + defer state.mu.Unlock() + return state.runEmitted +} + +// AgentMessageEventsSuppressedRun reports whether any Agent invocation +// suppressed terminal message emission during the entire canvas run. +func AgentMessageEventsSuppressedRun(ctx context.Context) bool { + state, ok := ctx.Value(agentMessageEmitterCtxKey{}).(*agentMessageEmitterState) + if !ok || state == nil { + return false + } + state.mu.Lock() + defer state.mu.Unlock() + return state.runSuppressed +} + +// EmitCanvasMessageEvent emits one presented Message event and falls back to +// the plain-content callback when no event callback is installed. +func EmitCanvasMessageEvent(ctx context.Context, content string, startToThink, endToThink bool) bool { + if emit, ok := ctx.Value(canvasMessageEventEmitterCtxKey{}).(CanvasMessageEventEmitter); ok && emit != nil { + emit(content, startToThink, endToThink) + if state, ok := ctx.Value(agentMessageEmitterCtxKey{}).(*agentMessageEmitterState); ok && state != nil { + state.mu.Lock() + state.emitted = true + state.runEmitted = true + state.mu.Unlock() + } + return true + } + if startToThink || endToThink { + return false + } + return EmitCanvasMessage(ctx, content) +} + +// EmitCanvasMessage emits already-rendered Message-component content through +// the direct canvas emitter, falling back to the Agent emitter when necessary. +// When the content exactly matches the answer already streamed by an upstream +// Agent, it is not emitted again. Distinct Message output is preserved. func EmitCanvasMessage(ctx context.Context, content string) bool { emit, ok := ctx.Value(canvasMessageEmitterCtxKey{}).(CanvasMessageEmitter) + state, hasAgentEmitter := ctx.Value(agentMessageEmitterCtxKey{}).(*agentMessageEmitterState) + if hasAgentEmitter && state != nil { + state.mu.Lock() + exactDuplicate := content != "" && state.agentContent.Len() > 0 && state.agentContent.String() == content + fallback := state.emit + state.mu.Unlock() + if exactDuplicate { + state.mu.Lock() + state.emitted = true + state.runEmitted = true + state.mu.Unlock() + return true + } + if ok && emit != nil { + emit(content) + state.mu.Lock() + if content != "" { + state.emitted = true + state.runEmitted = true + } + state.mu.Unlock() + return true + } + if fallback != nil { + fallback(content, "") + state.mu.Lock() + if content != "" { + state.emitted = true + state.runEmitted = true + } + state.mu.Unlock() + return true + } + return false + } if !ok || emit == nil { return false } emit(content) - if content != "" { - if state, ok := ctx.Value(agentMessageEmitterCtxKey{}).(*agentMessageEmitterState); ok && state != nil { - state.emitted = true - } - } return true } @@ -138,30 +383,81 @@ func EmitCanvasMessage(ctx context.Context, content string) bool { // 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 + if !ok || state == nil { + return false + } + state.mu.Lock() + defer state.mu.Unlock() + return state.emitted +} + +// DeferredAgentMessageEventsEmitted reports whether a deferred Agent has +// already delivered any deltas to its consuming Message node. It is kept +// separate from AgentMessageEventsEmitted so the model-stream collector can +// continue forwarding every delta instead of treating the first token as the +// entire response. +func DeferredAgentMessageEventsEmitted(ctx context.Context) bool { + sinkState := agentDeltaSinkFromContext(ctx) + if sinkState == nil { + return false + } + sinkState.mu.Lock() + defer sinkState.mu.Unlock() + return sinkState.emitted +} + +// HasDeferredAgentMessageSink reports whether the current invocation is being +// consumed by a downstream Message node. The model stream collector must not +// use the outer canvas-event flag for de-duplication in this mode: Message +// marks that flag after the first delta, while the remaining model deltas +// still need to reach the sink. +func HasDeferredAgentMessageSink(ctx context.Context) bool { + sinkState := agentDeltaSinkFromContext(ctx) + return sinkState != nil && sinkState.sink != nil } // 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 { + if !ok || state == nil { return } - if state.finalize() { + state.mu.Lock() + finalize := state.finalize + state.mu.Unlock() + if finalize == nil { + return + } + if finalize() { + state.mu.Lock() state.emitted = true + state.runEmitted = true + state.mu.Unlock() } } // ResetAgentMessageEmission starts a fresh Agent message emission scope. func ResetAgentMessageEmission(ctx context.Context) { + if sinkState := agentDeltaSinkFromContext(ctx); sinkState != nil { + sinkState.mu.Lock() + sinkState.emitted = false + sinkState.mu.Unlock() + } state, ok := ctx.Value(agentMessageEmitterCtxKey{}).(*agentMessageEmitterState) if !ok || state == nil { return } - if state.reset != nil { - state.reset() + state.mu.Lock() + reset := state.reset + state.mu.Unlock() + if reset != nil { + reset() } + state.mu.Lock() state.emitted = false + state.suppressed = false + state.agentContent.Reset() + state.mu.Unlock() } // GetStateFromContext extracts a typed state attached via WithState. diff --git a/internal/agent/runtime/context_test.go b/internal/agent/runtime/context_test.go new file mode 100644 index 0000000000..0523feb772 --- /dev/null +++ b/internal/agent/runtime/context_test.go @@ -0,0 +1,91 @@ +// +// 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 runtime + +import ( + "context" + "testing" + "time" +) + +func TestAgentMessageCallbacksRunOutsideStateLocks(t *testing.T) { + assertCompletes := func(name string, fn func()) { + t.Helper() + done := make(chan struct{}) + go func() { + fn() + close(done) + }() + select { + case <-done: + case <-time.After(time.Second): + t.Fatalf("%s callback re-entered a locked state", name) + } + } + + var agentCtx context.Context + agentCtx = WithAgentMessageEmitter(context.Background(), func(string, string) { + _ = AgentMessageEventsEmitted(agentCtx) + }) + assertCompletes("Agent emitter", func() { EmitAgentMessage(agentCtx, "token", "") }) + + var sinkCtx context.Context + sinkCtx = WithAgentDeltaSink(context.Background(), func(string, string) { + _ = DeferredAgentMessageEventsEmitted(sinkCtx) + }) + assertCompletes("deferred sink", func() { EmitAgentMessage(sinkCtx, "token", "") }) + + var canvasCtx context.Context + canvasCtx = WithAgentMessageEmitter(context.Background(), func(string, string) {}) + canvasCtx = WithCanvasMessageEmitter(canvasCtx, func(string) { + _ = AgentMessageEventsEmitted(canvasCtx) + }) + assertCompletes("Canvas emitter", func() { EmitCanvasMessage(canvasCtx, "message") }) + + var lifecycleCtx context.Context + lifecycleCtx = WithAgentMessageEmitterControl( + context.Background(), + func(string, string) {}, + func() bool { + _ = AgentMessageEventsEmitted(lifecycleCtx) + return false + }, + func() { _ = AgentMessageEventsEmitted(lifecycleCtx) }, + ) + assertCompletes("finalize", func() { FinalizeAgentMessage(lifecycleCtx) }) + assertCompletes("reset", func() { ResetAgentMessageEmission(lifecycleCtx) }) +} + +func TestAgentMessageRunStateSurvivesInvocationReset(t *testing.T) { + ctx := WithAgentMessageEmitter(context.Background(), func(string, string) {}) + EmitAgentMessage(ctx, "first", "") + ResetAgentMessageEmission(ctx) + + if AgentMessageEventsEmitted(ctx) { + t.Fatal("invocation emitted state survived reset") + } + if !AgentMessageEventsEmittedRun(ctx) { + t.Fatal("run emitted state was cleared by invocation reset") + } + + suppressedCtx := WithComponentExecutionOptions(ctx, ComponentExecutionOptions{SuppressAgentMessageEvents: true}) + EmitAgentMessage(suppressedCtx, "hidden", "") + ResetAgentMessageEmission(ctx) + if !AgentMessageEventsSuppressedRun(ctx) { + t.Fatal("run suppressed state was cleared by invocation reset") + } +} diff --git a/internal/entity/models/llm.go b/internal/entity/models/llm.go index 985e3c917c..13bd54cfb9 100644 --- a/internal/entity/models/llm.go +++ b/internal/entity/models/llm.go @@ -303,9 +303,6 @@ func (m *EinoChatModel) Stream(ctx context.Context, msgs []*schema.Message, opts 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") } diff --git a/internal/entity/models/llm_test.go b/internal/entity/models/llm_test.go index b94bbfc957..7f00028c76 100644 --- a/internal/entity/models/llm_test.go +++ b/internal/entity/models/llm_test.go @@ -13,15 +13,8 @@ import ( func TestEinoChatModelStreamFiltersDoneSentinel(t *testing.T) { modelName := "chat" driver := &streamSentinelDriver{captureToolDriver: &captureToolDriver{}} - var callbacks []string base := NewChatModel(driver, &modelName, &APIConfig{}) - model := NewEinoChatModel(base, &ChatConfig{ - StreamCallback: func(content, _ string) { - if content != "" { - callbacks = append(callbacks, content) - } - }, - }) + model := NewEinoChatModel(base, &ChatConfig{}) stream, err := model.Stream(context.Background(), []*schema.Message{schema.UserMessage("hello")}) if err != nil { @@ -45,9 +38,6 @@ func TestEinoChatModelStreamFiltersDoneSentinel(t *testing.T) { if len(messages) != 2 || messages[0] != "answer" || messages[1] != "DONE!" { t.Fatalf("stream messages = %#v, want [answer DONE!]", messages) } - if len(callbacks) != 2 || callbacks[0] != "answer" || callbacks[1] != "DONE!" { - t.Fatalf("stream callbacks = %#v, want [answer DONE!]", callbacks) - } } func TestEinoChatModelGenerateSendsBoundTools(t *testing.T) { @@ -166,15 +156,8 @@ func TestEinoChatModelStreamWithToolsStreamsFinalAnswer(t *testing.T) { driver := &captureToolDriver{ resp: &ChatResponse{Answer: &answer}, } - var callbacks []string base := NewChatModel(driver, &modelName, &APIConfig{ApiKey: &apiKey}) - model := NewEinoChatModel(base, &ChatConfig{ - StreamCallback: func(content, _ string) { - if content != "" { - callbacks = append(callbacks, content) - } - }, - }) + model := NewEinoChatModel(base, &ChatConfig{}) bound, err := model.WithTools([]*schema.ToolInfo{ { Name: "search_my_dateset", @@ -206,9 +189,6 @@ func TestEinoChatModelStreamWithToolsStreamsFinalAnswer(t *testing.T) { if msg == nil || msg.Content != answer { t.Fatalf("stream message = %#v, want final answer content", msg) } - if len(callbacks) != 1 || callbacks[0] != answer { - t.Fatalf("callbacks = %#v, want streamed answer", callbacks) - } } func TestToInternalMessagesPreservesToolMessages(t *testing.T) { diff --git a/internal/entity/models/types.go b/internal/entity/models/types.go index a16864243a..e5b4c09345 100644 --- a/internal/entity/models/types.go +++ b/internal/entity/models/types.go @@ -182,9 +182,6 @@ type ChatConfig struct { // non-nil) after the stream completes; callers read it the same // way they read ToolCallsResult. UsageResult *TokenUsage `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/service/agent.go b/internal/service/agent.go index 9d08ee4f2c..64e7b6b7a9 100644 --- a/internal/service/agent.go +++ b/internal/service/agent.go @@ -1242,11 +1242,19 @@ func (s *AgentService) buildRunFunc(canvasID string, versionRow *entity.UserCanv TaskID: taskID, SessionID: sessionID, }) + ctx2 = runtime.WithDeferredNodeRegistry(ctx2) agentMessageEmit, agentMessageFinalize, agentMessageReset := makeAgentMessageDeltaEmitterWithFinalizer(emit) ctx2 = runtime.WithAgentMessageEmitterControl(ctx2, agentMessageEmit, agentMessageFinalize, agentMessageReset) ctx2 = runtime.WithCanvasMessageEmitter(ctx2, func(content string) { emitAgentMessageEvent(emit, canvas.MessageEvent{Content: content}) }) + ctx2 = runtime.WithCanvasMessageEventEmitter(ctx2, func(content string, startToThink, endToThink bool) { + emitAgentMessageEvent(emit, canvas.MessageEvent{ + Content: content, + StartToThink: startToThink, + EndToThink: endToThink, + }) + }) // Seed initial env/sys values from the Canvas DSL globals. // Python's self.globals dict stores "sys.*" and "env.*" under @@ -1410,8 +1418,13 @@ func (s *AgentService) buildRunFunc(canvasID string, versionRow *entity.UserCanv } referencePayload := agentRunReferencePayload(state, legacyReference) assistantOutput := terminalCanvasOutput(c, state, workflowOutput, answer, downloads) + // Release any deferred Agent node that was not consumed because the + // downstream Message was skipped by an exception/branch path. + runtime.CompleteAllDeferredNodes(ctx2) runtime.FinalizeAgentMessage(ctx2) - messageEventsEmitted := runtime.AgentMessageEventsEmitted(ctx2) + messageEventsEmitted := runtime.AgentMessageEventsEmittedRun(ctx2) + messageEventsSuppressed := runtime.AgentMessageEventsSuppressedRun(ctx2) + shouldEmitMessage := messageEventsEmitted || !messageEventsSuppressed if err != nil { common.Debug("RunAgent invoke err", @@ -1429,7 +1442,7 @@ func (s *AgentService) buildRunFunc(canvasID string, versionRow *entity.UserCanv if persistErr := s.persistAgentRunSession(ctx, canvasID, userID, sessionID, messageID, userInput, answer, referencePayload, dsl, state, answer != ""); persistErr != nil { return nil, fmt.Errorf("persist interrupted agent session: %w: %w", persistErr, ErrAgentStorageError) } - if answer != "" { + if answer != "" && shouldEmitMessage { if !messageEventsEmitted { emitAgentMessageEvents(emit, answer, thinking, referencePayload) } @@ -1447,14 +1460,16 @@ func (s *AgentService) buildRunFunc(canvasID string, versionRow *entity.UserCanv s.markRunFailed(ctx2, runID, "persist session: "+persistErr.Error()) return nil, fmt.Errorf("persist agent session: %w: %w", persistErr, ErrAgentStorageError) } - if !messageEventsEmitted { + if !messageEventsEmitted && shouldEmitMessage { emitAgentMessageEvents(emit, answer, thinking, referencePayload) } - meData, _ := json.Marshal(canvas.MessageEndEvent{ - Reference: referencePayload, - }) - emit("message_end", string(meData)) + if shouldEmitMessage { + meData, _ := json.Marshal(canvas.MessageEndEvent{ + Reference: referencePayload, + }) + emit("message_end", string(meData)) + } wfPayload := map[string]interface{}{ "inputs": map[string]any{"query": userInput}, @@ -1481,14 +1496,16 @@ func (s *AgentService) buildRunFunc(canvasID string, versionRow *entity.UserCanv s.markRunFailed(ctx2, runID, "persist session: "+persistErr.Error()) return nil, fmt.Errorf("persist agent session: %w: %w", persistErr, ErrAgentStorageError) } - if !messageEventsEmitted { + if !messageEventsEmitted && shouldEmitMessage { emitAgentMessageEvents(emit, answer, thinking, referencePayload) } - meData, _ := json.Marshal(canvas.MessageEndEvent{ - Reference: referencePayload, - }) - emit("message_end", string(meData)) + if shouldEmitMessage { + meData, _ := json.Marshal(canvas.MessageEndEvent{ + Reference: referencePayload, + }) + emit("message_end", string(meData)) + } // Emit workflow_finished with the final outputs and aggregated // per-run token usage across all LLM calls in this turn.