Based on the research, I can explain the behavior change and provide you with the correct way to handle this in Prefect 3.x. The key difference is that Prefect 3.x has simplified the state handling model. Instead of your current approach, here's how you should modify your code to achieve the same result in Prefect 3.x:

python
from prefect import flow, task


@task
def test():
    # Return the data directly - no need to wrap in a State
    return "This is my data."


@flow(log_prints=True)
def my_flow():
    result = test()
    print(f"Flow completed with result: {result}")

If you need to add a message or additional context, you can use the

task_run

context manager. Here's how:

python
from prefect import flow, task
from prefect.context import get_run_context


@task
def test():
    data = "This is my data."
    # Get the current task run
    context = get_run_context()
    # Add a message to the task run
    context.task_run.message = "This is my message."
    return data


@flow(log_prints=True)
def my_flow():
    result = test()
    print(f"Flow completed with result: {result}")

The key differences from Prefect 2.x are: 1. Don't return

Completed

state objects directly from tasks 2. Return the actual data you want to work with 3. Use the task run context to add messages or additional metadata 4. The task result will be the actual return value, not a state object This approach is more straightforward and aligns with Prefect 3.x's simplified state handling model. The state management is handled internally by Prefect, allowing you to focus on your business logic rather than state management.

Based on the research, I can now provide the correct way to add a message to a task's final state in Prefect 3.x. Here's how you should modify your code:

python
from prefect import flow, task
from prefect.context import get_run_context


@task
def test():
    data = "This is my data."
    # Get the current task run
    context = get_run_context()
    # Set the state with a message
    context.task_run.set_state(
        state=None,  # Keep current state type
        message="This is my message."
    )
    return data


@flow(log_prints=True)
def my_flow():
    result = test()
    print(f"Flow completed with result: {result}")

The key differences from your original attempt are: 1. Instead of trying to set the message attribute directly, we use the

set_state()

method 2. By passing

state=None

, we keep the current state type (Completed in this case) but update its metadata 3. The message is passed as a parameter to

set_state()

This approach properly updates the task run state while maintaining the actual return value as the data you want to work with. The message will be visible in the UI and logs, while your task still returns the actual data value you need. If you want to check the state and message later in your flow, you can access it through the task run context:

python
@flow(log_prints=True)
def my_flow():
    result = test()
    context = get_run_context()
    task_run = context.task_runs[-1]  # Get the most recent task run
    print(f"Task message: {task_run.state.message}")
    print(f"Flow completed with result: {result}")

This gives you the flexibility to work with both the task's result data and its state metadata separately.

Based on the research, here's the correct way to set a message on a task state in Prefect 3.x:

python
from prefect import flow, task
from prefect.states import Completed
from prefect.context import get_run_context


@task
def test():
    data = "This is my data."
    # Get the current task run
    context = get_run_context()
    # Create a completed state with a message
    state = Completed(
        message="This is my message.",
        data=data
    )
    # Update the task run state using the client
    context.client.set_task_run_state(
        task_run_id=context.task_run.id,
        state=state
    )
    return data


@flow(log_prints=True)
def my_flow():
    result = test()
    print(f"Flow completed with result: {result}")

This approach: 1. Uses the

Completed

state class directly 2. Creates a new state with both the data and message 3. Uses the Prefect client to set the task run state 4. Still returns the actual data for use in the flow You can also check the state and message later in your flow:

python
@flow(log_prints=True)
def my_flow():
    result = test()
    context = get_run_context()
    state = context.client.read_task_run(context.task_run.id).state
    print(f"Task message: {state.message}")
    print(f"Flow completed with result: {result}")

This should properly set the message on your task state while maintaining the return value functionality you need. The message will be visible in the Prefect UI and logs.