def custom_function_step(step_input: StepInput, session_state):
    session_state["test"] = "test_1"  # Write to session state
    return StepOutput(content=f"Updated session_state: {session_state}")

        __
        __