diff --git a/src/ppytty/kernel/common.py b/src/ppytty/kernel/common.py
index 2796add..a730fcb 100644
--- a/src/ppytty/kernel/common.py
+++ b/src/ppytty/kernel/common.py
@@ -11,7 +11,9 @@ from . state import state
 
 def runnable_task(task):
 
-    return task() if callable(task) else task
+    runnable = task() if callable(task) else task
+    state.spawned_objects[runnable] = task
+    return runnable
 
 
 
diff --git a/src/ppytty/kernel/scheduler.py b/src/ppytty/kernel/scheduler.py
index a392f8b..b4e6052 100644
--- a/src/ppytty/kernel/scheduler.py
+++ b/src/ppytty/kernel/scheduler.py
@@ -162,7 +162,9 @@ def process_task_completion(task, success, result):
     if not candidate_parent and task is not state.top_task:
         log.error('%r completed with no parent', task)
     if candidate_parent in state.tasks_waiting_child:
-        common.trap_will_return(candidate_parent, (task, success, result))
+        spawned_object = state.spawned_objects[task]
+        common.trap_will_return(candidate_parent, (spawned_object, success, result))
+        del state.spawned_objects[task]
         del state.parent_task[task]
         state.child_tasks[candidate_parent].remove(task)
         common.clear_tasks_children(candidate_parent)
diff --git a/src/ppytty/kernel/state.py b/src/ppytty/kernel/state.py
index ac6aee0..824c939 100644
--- a/src/ppytty/kernel/state.py
+++ b/src/ppytty/kernel/state.py
@@ -51,6 +51,12 @@ class _State(object):
         self.tasks_waiting_time = []
         self.tasks_waiting_time_hq = []
 
+        # ---------------------------------------------------------------------
+        # Spawned objects
+
+        # Maps running tasks to objects passed to scheduler.run / task-spawn.
+        self.spawned_objects = {}
+
         # ---------------------------------------------------------------------
         # Task trap tracking.
 
diff --git a/src/ppytty/kernel/trap_handlers.py b/src/ppytty/kernel/trap_handlers.py
index f367e08..f148354 100644
--- a/src/ppytty/kernel/trap_handlers.py
+++ b/src/ppytty/kernel/trap_handlers.py
@@ -139,7 +139,9 @@ def task_wait(task):
             success, result = state.completed_tasks[child]
             break
     if child is not None:
-        common.trap_will_return(task, (child, success, result))
+        spawned_object = state.spawned_objects[child]
+        common.trap_will_return(task, (spawned_object, success, result))
+        del state.spawned_objects[child]
         del state.parent_task[child]
         state.child_tasks[task].remove(child)
         common.clear_tasks_children(task)
@@ -198,7 +200,8 @@ def message_send(task, to_task, message):
 
     if to_task in state.tasks_waiting_inbox:
         state.tasks_waiting_inbox.remove(to_task)
-        common.trap_will_return(to_task, (task, message))
+        spawned_object = state.spawned_objects[task]
+        common.trap_will_return(to_task, (spawned_object, message))
         state.runnable_tasks.append(to_task)
     else:
         state.task_inbox[to_task].append((task, message))
@@ -212,7 +215,8 @@ def message_wait(task):
     task_inbox = state.task_inbox[task]
     if task_inbox:
         sender_task, message = task_inbox.popleft()
-        common.trap_will_return(task, (sender_task, message))
+        spawned_object = state.spawned_objects[sender_task]
+        common.trap_will_return(task, (spawned_object, message))
         state.runnable_tasks.append(task)
     else:
         state.tasks_waiting_inbox.append(task)

Traps returning tasks need fixing.
Summary:
* Traps like `task-wait` and `message-wait` return, among other things, a reference to a task.
* The returned value is, however, not necessarily the same that was used in `task-spawn`.

Motive:
* The kernel will call the passed in task object if it is callable (assuming it is a generator function) which is then used in those traps results.

Failing example:
```
def child():
    yield ('sleep',)

def parent():
    yield ('task-spawn', child)
    completed_task, success, result = yield ('task-wait',)
    # This fails because completed_task is actually the generator object obtained by calling 
    assert completed_task is childchild
```

Possible solutions:
* Have the kernel require passed in tasks (to top-level `run`/`task-spawn`) to be a *generator object* or alike.
* Have the kernel track what it was passed vs. what it uses internally.

Let's sleep on this for a while.
diff --git a/tests/kernel/tasks.py b/tests/kernel/tasks.py
index c66aa2e..2397d6b 100644
--- a/tests/kernel/tasks.py
+++ b/tests/kernel/tasks.py
@@ -13,11 +13,13 @@ def sleep_zero():
 
 
 
-def spawn_wait(task):
+def spawn_wait(task, sleep_before_wait=False):
 
     yield ('task-spawn', task)
-    _, success, result = yield ('task-wait',)
-    return success, result
+    if sleep_before_wait:
+        yield ('sleep', 0)
+    completed_task, success, result = yield ('task-wait',)
+    return completed_task, success, result
 
 
 
diff --git a/tests/kernel/test_state.py b/tests/kernel/test_state.py
index 0a66aa6..6537cb3 100644
--- a/tests/kernel/test_state.py
+++ b/tests/kernel/test_state.py
@@ -79,6 +79,7 @@ STATE_ATTRS = {
     'tasks_waiting_key_hq': (_assert_empty_list, _change_list),
     'tasks_waiting_time': (_assert_empty_list, _change_list),
     'tasks_waiting_time_hq': (_assert_empty_list, _change_list),
+    'spawned_objects': (_assert_empty_dict, _change_dict),
     'trap_call': (_assert_empty_dict, _change_dict),
     'trap_success': (_assert_empty_dict, _change_dict),
     'trap_result': (_assert_empty_dict, _change_dict),
diff --git a/tests/kernel/test_task_api_msgs.py b/tests/kernel/test_task_api_msgs.py
index ccb15bc..adfc29b 100644
--- a/tests/kernel/test_task_api_msgs.py
+++ b/tests/kernel/test_task_api_msgs.py
@@ -119,4 +119,72 @@ class TestChildToParent(io_bypass.NoOutputTestCase):
         self.assertEqual(result['received-message'], 'child-to-parent-message')
 
 
+
+class TestSenderTaskIsChildTask(io_bypass.NoOutputTestCase):
+
+    def _fast_child(self):
+
+        yield ('message-send', None, 'hi-from-child')
+
+
+    def _slow_child(self):
+
+        yield ('sleep', 0)
+        yield ('message-send', None, 'hi-from-child')
+
+
+    def _parent(self, child_object):
+
+        yield ('task-spawn', child_object)
+        sender_task, message = yield ('message-wait',)
+        completed_child, child_success, child_result = yield ('task-wait',)
+        return sender_task, message, completed_child, child_success, child_result
+
+
+    def _assert_parent_result(self, child_object, result):
+
+        sender_task, message, completed_child, child_success, child_result = result
+        self.assertIs(sender_task, child_object)
+        self.assertEqual(message, 'hi-from-child')
+        self.assertIs(completed_child, child_object)
+        self.assertTrue(child_success)
+        self.assertIsNone(child_result)
+
+
+    def test_gen_function_child_parent_waits_1st(self):
+
+        child_gen_function = self._slow_child
+        parent = self._parent(child_gen_function)
+        success, result = run(parent)
+        self.assertTrue(success)
+        self._assert_parent_result(child_gen_function, result)
+
+
+    def test_gen_function_child_parent_waits_2nd(self):
+
+        child_gen_function = self._fast_child
+        parent = self._parent(child_gen_function)
+        success, result = run(parent)
+        self.assertTrue(success)
+        self._assert_parent_result(child_gen_function, result)
+
+
+    def test_gen_object_child_parent_waits_1st(self):
+
+        child_gen_object = self._slow_child()
+        parent = self._parent(child_gen_object)
+        success, result = run(parent)
+        self.assertTrue(success)
+        self._assert_parent_result(child_gen_object, result)
+
+
+    def test_gen_object_child_parent_waits_2nd(self):
+
+        child_gen_object = self._fast_child()
+        parent = self._parent(child_gen_object)
+        success, result = run(parent)
+        self.assertTrue(success)
+        self._assert_parent_result(child_gen_object, result)
+
+
 # ----------------------------------------------------------------------------
diff --git a/tests/kernel/test_task_api_tasks.py b/tests/kernel/test_task_api_tasks.py
index d26946e..9893e37 100644
--- a/tests/kernel/test_task_api_tasks.py
+++ b/tests/kernel/test_task_api_tasks.py
@@ -14,20 +14,54 @@ from . import tasks
 
 
 
-class TestSpawn(io_bypass.NoOutputTestCase):
+class TestSpawnWaitObjects(io_bypass.NoOutputTestCase):
 
-    def test_spawn_gen_function(self):
+    def test_spawn_wait_gen_function_child_completes_1st(self):
 
         generator_function = tasks.sleep_zero
         task = tasks.spawn_wait(generator_function)
-        run(task)
+        success, result = run(task)
+        self.assertTrue(success)
+        completed_child, child_success, child_result = result
+        self.assertIs(completed_child, generator_function)
+        self.assertTrue(child_success)
+        self.assertIsNone(child_result)
 
 
-    def test_spawn_gen_object(self):
+    def test_spawn_wait_gen_object_child_completes_1st(self):
 
         generator_object = tasks.sleep_zero()
         task = tasks.spawn_wait(generator_object)
-        run(task)
+        success, result = run(task)
+        self.assertTrue(success)
+        completed_child, child_success, child_result = result
+        self.assertIs(completed_child, generator_object)
+        self.assertTrue(child_success)
+        self.assertIsNone(child_result)
+
+
+    def test_spawn_wait_gen_function_child_completes_2nd(self):
+
+        generator_function = tasks.sleep_zero
+        task = tasks.spawn_wait(generator_function, sleep_before_wait=True)
+        success, result = run(task)
+        self.assertTrue(success)
+        completed_child, child_success, child_result = result
+        self.assertIs(completed_child, generator_function)
+        self.assertTrue(child_success)
+        self.assertIsNone(child_result)
+
+
+    def test_spawn_wait_gen_object_child_completes_2nd(self):
+
+        generator_object = tasks.sleep_zero()
+        task = tasks.spawn_wait(generator_object, sleep_before_wait=True)
+        success, result = run(task)
+        self.assertTrue(success)
+        completed_child, child_success, child_result = result
+        self.assertIs(completed_child, generator_object)
+        self.assertTrue(child_success)
+        self.assertIsNone(child_result)
 
 
 
@@ -38,7 +72,7 @@ class TestWait(io_bypass.NoOutputTestCase):
         task = tasks.spawn_wait(tasks.sleep_zero_return_42_idiv_arg(42))
         parent_sucess, parent_result = run(task)
         self.assertTrue(parent_sucess)
-        child_success, child_result = parent_result
+        _completed_child, child_success, child_result = parent_result
         self.assertTrue(child_success)
         self.assertEqual(child_result, 1)
 
@@ -48,7 +82,7 @@ class TestWait(io_bypass.NoOutputTestCase):
         task = tasks.spawn_wait(tasks.sleep_zero_return_42_idiv_arg(0))
         parent_success, parent_result = run(task)
         self.assertTrue(parent_success)
-        child_success, child_result = parent_result
+        _completed_child, child_success, child_result = parent_result
         self.assertFalse(child_success)
         self.assertIsInstance(child_result, ZeroDivisionError)
 
@@ -58,15 +92,10 @@ class TestWaitChildException(io_bypass.NoOutputTestCase):
 
     def parent_spawn_wait_child_exception(self, child_task, expected_exc_class):
 
-        def parent():
-            yield ('task-spawn', child_task)
-            completed_task, child_success, child_result = yield ('task-wait',)
-            return completed_task is child_task, child_success, child_result
-
-        success, result = run(parent)
+        success, result = run(tasks.spawn_wait(child_task))
         self.assertTrue(success)
-        correct_task_completed, child_task_success, child_task_result = result
-        self.assertTrue(correct_task_completed)
+        completed_child, child_task_success, child_task_result = result
+        self.assertIs(completed_child, child_task)
         self.assertFalse(child_task_success)
         self.assertIsInstance(child_task_result, expected_exc_class)
 

pip install -e .[dev]
parse_log_pytest
requirements.txt
pytest
apt-get update apt-get install -y gcc
3.9
requirements/base.txt
pytest --no-header -rA --tb=line --color=no -p no:cacheprovider -W ignore::DeprecationWarning
astroid==3.3.9
blessings==1.7
coverage==7.8.0
dill==0.3.9
exceptiongroup==1.2.2
iniconfig==2.1.0
isort==6.0.1
mccabe==0.7.0
packaging==24.2
platformdirs==4.3.7
pluggy==1.5.0
-e git+https://github.com/tmontes/ppytty.git@36a1a711020ff143e761f203127ccebc627b85c5#egg=ppytty
pylint==3.3.6
pyte==0.8.2
pytest==8.3.5
six==1.17.0
tomli==2.2.1
tomlkit==0.13.2
typing_extensions==4.13.0
wcwidth==0.2.13

name: ppytty
channels:
  - defaults
  - https://repo.anaconda.com/pkgs/main
  - https://repo.anaconda.com/pkgs/r
  - conda-forge
dependencies:
  - _libgcc_mutex=0.1=main
  - _openmp_mutex=5.1=1_gnu
  - ca-certificates=2025.2.25=h06a4308_0
  - ld_impl_linux-64=2.40=h12ee557_0
  - libffi=3.4.4=h6a678d5_1
  - libgcc-ng=11.2.0=h1234567_1
  - libgomp=11.2.0=h1234567_1
  - libstdcxx-ng=11.2.0=h1234567_1
  - ncurses=6.4=h6a678d5_0
  - openssl=3.0.16=h5eee18b_0
  - pip=25.0=py39h06a4308_0
  - python=3.9.21=he870216_1
  - readline=8.2=h5eee18b_0
  - setuptools=75.8.0=py39h06a4308_0
  - sqlite=3.45.3=h5eee18b_0
  - tk=8.6.14=h39e8969_0
  - tzdata=2025a=h04d1e81_0
  - wheel=0.45.1=py39h06a4308_0
  - xz=5.6.4=h5eee18b_1
  - zlib=1.2.13=h5eee18b_1
  - pip:
      - astroid==3.3.9
      - blessings==1.7
      - coverage==7.8.0
      - dill==0.3.9
      - exceptiongroup==1.2.2
      - iniconfig==2.1.0
      - isort==6.0.1
      - mccabe==0.7.0
      - packaging==24.2
      - platformdirs==4.3.7
      - pluggy==1.5.0
      - pylint==3.3.6
      - pyte==0.8.2
      - pytest==8.3.5
      - six==1.17.0
      - tomli==2.2.1
      - tomlkit==0.13.2
      - typing-extensions==4.13.0
      - wcwidth==0.2.13
prefix: /opt/conda/envs/ppytty

tests/kernel/test_state.py::TestRun::test_state_attrs_are_clean tests/kernel/test_state.py::TestRun::test_state_attrs_are_these tests/kernel/test_state.py::TestRun::test_state_reset
tests/kernel/test_task_api_msgs.py::TestParentToChild::test_parent_spawn_child_message_wait_parent_send tests/kernel/test_task_api_msgs.py::TestParentToChild::test_parent_spawn_send_child_message_wait tests/kernel/test_task_api_msgs.py::TestChildToParent::test_parent_spawn_child_message_send_parent_message_wait tests/kernel/test_task_api_msgs.py::TestChildToParent::test_parent_spawn_message_wait_child_message_send tests/kernel/test_task_api_msgs.py::TestSenderTaskIsChildTask::test_gen_function_child_parent_waits_1st tests/kernel/test_task_api_msgs.py::TestSenderTaskIsChildTask::test_gen_function_child_parent_waits_2nd tests/kernel/test_task_api_msgs.py::TestSenderTaskIsChildTask::test_gen_object_child_parent_waits_1st tests/kernel/test_task_api_msgs.py::TestSenderTaskIsChildTask::test_gen_object_child_parent_waits_2nd tests/kernel/test_task_api_tasks.py::TestSpawnWaitObjects::test_spawn_wait_gen_function_child_completes_1st tests/kernel/test_task_api_tasks.py::TestSpawnWaitObjects::test_spawn_wait_gen_function_child_completes_2nd tests/kernel/test_task_api_tasks.py::TestSpawnWaitObjects::test_spawn_wait_gen_object_child_completes_1st tests/kernel/test_task_api_tasks.py::TestSpawnWaitObjects::test_spawn_wait_gen_object_child_completes_2nd tests/kernel/test_task_api_tasks.py::TestWait::test_wait_child_exception tests/kernel/test_task_api_tasks.py::TestWait::test_wait_child_success tests/kernel/test_task_api_tasks.py::TestWaitChildException::test_parent_spawn_wait_child_exception tests/kernel/test_task_api_tasks.py::TestWaitChildException::test_parent_spawn_wait_child_trap_arg_count_wrong_exception tests/kernel/test_task_api_tasks.py::TestWaitChildException::test_parent_spawn_wait_child_trap_does_not_exist_exception tests/kernel/test_task_api_tasks.py::TestSpawnDontWait::test_child_spawn_crash_no_grand_child_wait_child_completes_1st tests/kernel/test_task_api_tasks.py::TestSpawnDontWait::test_child_spawn_crash_no_grand_child_wait_child_completes_2nd tests/kernel/test_task_api_tasks.py::TestSpawnDontWait::test_parent_spawn_crash_no_child_wait_parent_completes_1st tests/kernel/test_task_api_tasks.py::TestSpawnDontWait::test_parent_spawn_crash_no_child_wait_parent_completes_2nd
src/ppytty/kernel/common.py src/ppytty/kernel/trap_handlers.py src/ppytty/kernel/state.py src/ppytty/kernel/scheduler.py
src/ppytty/kernel/common.py src/ppytty/kernel/trap_handlers.py src/ppytty/kernel/state.py src/ppytty/kernel/scheduler.py