fix: failed to set right status in memory (#17472)

As title

Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>

internal/agent/component/message.go

@@ -316,8 +316,8 @@ func (m *MessageComponent) Invoke(ctx context.Context, db *gorm.DB, inputs map[s
 		saveErr := saver.Save(ctx, MemorySaveRequest{
 			MemoryIDs:     memIDs,
 			UserID:        userID,
-			AgentID:       stringFromStateSys(state, "agent_id"),
-			SessionID:     state.SessionID,
+			AgentID:       memoryAgentID(state),
+			SessionID:     memorySessionID(state),
 			UserInput:     stringFromStateSys(state, "query"),
 			AgentResponse: rendered,
 		})
@@ -330,6 +330,29 @@ func (m *MessageComponent) Invoke(ctx context.Context, db *gorm.DB, inputs map[s
 	return out, nil
 }
 
+func memoryAgentID(state *runtime.CanvasState) string {
+	if agentID := stringFromStateSys(state, "agent_id"); agentID != "" {
+		return agentID
+	}
+	if canvasID := stringFromStateSys(state, "canvas_id"); canvasID != "" {
+		return canvasID
+	}
+	if state == nil {
+		return ""
+	}
+	return state.SessionID
+}
+
+func memorySessionID(state *runtime.CanvasState) string {
+	if sessionID := stringFromStateSys(state, "session_id"); sessionID != "" {
+		return sessionID
+	}
+	if state == nil {
+		return ""
+	}
+	return state.RunID
+}
+
 // 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.

internal/agent/component/message_phase8b_test.go

@@ -257,6 +257,8 @@ func TestMessage_MemorySave_Success(t *testing.T) {
 	c, _ := NewMessageComponent(map[string]any{"text": "hi"})
 	state := canvas.NewCanvasState("run-y", "session-y")
 	state.Sys["query"] = "what?"
+	state.Sys["canvas_id"] = "canvas-y"
+	state.Sys["session_id"] = "session-y"
 	state.Sys["agent_id"] = "agent-y"
 	ctx := withStateForTest(context.Background(), state)
 
@@ -286,6 +288,37 @@ func TestMessage_MemorySave_Success(t *testing.T) {
 	}
 }
 
+func TestMessage_MemorySave_FallbackIDs(t *testing.T) {
+	var saved MemorySaveRequest
+	saveFn := memSaverFunc(func(_ context.Context, req MemorySaveRequest) error {
+		saved = req
+		return nil
+	})
+	SetMemorySaver(&saveFn)
+	defer SetMemorySaver(nil)
+
+	c, _ := NewMessageComponent(map[string]any{"text": "hi"})
+	state := canvas.NewCanvasState("run-fallback", "task-fallback")
+	state.Sys["query"] = "what?"
+	ctx := withStateForTest(context.Background(), state)
+
+	_, err := c.Invoke(ctx, nil, map[string]any{
+		"text":        "hi",
+		"stream":      false,
+		"memory_save": true,
+		"memory_ids":  []string{"m1"},
+	})
+	if err != nil {
+		t.Fatalf("Invoke: %v", err)
+	}
+	if saved.AgentID != "task-fallback" {
+		t.Errorf("AgentID fallback: got %q, want task-fallback", saved.AgentID)
+	}
+	if saved.SessionID != "run-fallback" {
+		t.Errorf("SessionID fallback: got %q, want run-fallback", saved.SessionID)
+	}
+}
+
 // TestMessage_MemorySave_FromDSLParams: memory_ids declared in the
 // DSL params (constructor) are honoured even when the runtime inputs
 // map does not carry memory_ids. This is the production path where