diff --git a/.changeset/khaki-emus-jog.md b/.changeset/khaki-emus-jog.md
new file mode 100644
index 0000000..6a38cf8
--- /dev/null
+++ b/.changeset/khaki-emus-jog.md
@@ -0,0 +1,26 @@
+---
+'@statelyai/agent': minor
+---
+
+Correlation IDs are now provided as part of the result from `agent.generateText(…)` and `agent.streamText(…)`:
+
+```ts
+const result = await agent.generateText({
+  prompt: 'Write me a song',
+  correlationId: 'my-correlation-id',
+  // ...
+});
+
+result.correlationId; // 'my-correlation-id'
+```
+
+These correlation IDs can be passed to feedback:
+
+```ts
+// ...
+
+agent.addFeedback({
+  reward: -1,
+  correlationId: result.correlationId,
+});
+```
diff --git a/.changeset/silent-deers-visit.md b/.changeset/silent-deers-visit.md
new file mode 100644
index 0000000..d557210
--- /dev/null
+++ b/.changeset/silent-deers-visit.md
@@ -0,0 +1,11 @@
+---
+'@statelyai/agent': minor
+---
+
+Changes to agent feedback (the `AgentFeedback` interface):
+
+- `goal` is now optional
+- `observationId` is now optional
+- `correlationId` has been added (optional)
+- `reward` has been added (optional)
+- `attributes` are now optional
diff --git a/package.json b/package.json
index 85733e4..c888ec2 100644
--- a/package.json
+++ b/package.json
@@ -38,7 +38,7 @@
     "ts-node": "^10.9.2",
     "tsup": "^8.2.0",
     "typescript": "^5.5.3",
-    "vitest": "^1.6.0",
+    "vitest": "^2.0.3",
     "wikipedia": "^2.1.2",
     "zod": "^3.23.8"
   },
diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml
index eb13c7e..515639d 100644
--- a/pnpm-lock.yaml
+++ b/pnpm-lock.yaml
@@ -59,8 +59,8 @@ devDependencies:
     specifier: ^5.5.3
     version: 5.5.3
   vitest:
-    specifier: ^1.6.0
-    version: 1.6.0(@types/node@20.14.11)
+    specifier: ^2.0.3
+    version: 2.0.3(@types/node@20.14.11)
   wikipedia:
     specifier: ^2.1.2
     version: 2.1.2
@@ -230,7 +230,6 @@ packages:
     dependencies:
       '@jridgewell/gen-mapping': 0.3.5
       '@jridgewell/trace-mapping': 0.3.25
-    dev: false
 
   /@babel/helper-string-parser@7.24.8:
     resolution: {integrity: sha512-pO9KhhRcuUyGnJWwyEgnRJTSIZHiT+vMD0kPeD+so0l7mxkMT19g3pjY9GTnHySck/hDzq+dtW/4VgnMkippsQ==, tarball: https://registry.npmjs.org/@babel/helper-string-parser/-/helper-string-parser-7.24.8.tgz}
@@ -920,13 +919,6 @@ packages:
       wrap-ansi-cjs: /wrap-ansi@7.0.0
     dev: true
 
-  /@jest/schemas@29.6.3:
-    resolution: {integrity: sha512-mo5j5X+jIZmJQveBKeS/clAueipV7KgiX1vMgCxam1RNYiqE1w62n0/tJJnHtjW8ZHcQco5gY85jA3mi0L+nSA==, tarball: https://registry.npmjs.org/@jest/schemas/-/schemas-29.6.3.tgz}
-    engines: {node: ^14.15.0 || ^16.10.0 || >=18.0.0}
-    dependencies:
-      '@sinclair/typebox': 0.27.8
-    dev: true
-
   /@jridgewell/gen-mapping@0.3.5:
     resolution: {integrity: sha512-IzL8ZoEDIBRWEzlCcRhOaCupYyN5gdIK+Q6fbFdPDg6HqX6jpkItn7DFIpW9LQzXG6Df9sA7+OKnq0qlz/GaQg==, tarball: https://registry.npmjs.org/@jridgewell/gen-mapping/-/gen-mapping-0.3.5.tgz}
     engines: {node: '>=6.0.0'}
@@ -1492,10 +1484,6 @@ packages:
     dev: true
     optional: true
 
-  /@sinclair/typebox@0.27.8:
-    resolution: {integrity: sha512-+Fj43pSMwJs4KRrH/938Uf+uAELIgVBmQzg/q1YG10djyfA3TnrU8N8XzqCh/okZdszqBQTZf96idMfE5lnwTA==, tarball: https://registry.npmjs.org/@sinclair/typebox/-/typebox-0.27.8.tgz}
-    dev: true
-
   /@tsconfig/node10@1.0.11:
     resolution: {integrity: sha512-DcRjDCujK/kCk/cUe8Xz8ZSpm8mS3mNNpta+jGCA6USEDfktlNvm1+IuZ9eTcDbNk41BHwpHHeW+N1lKCz4zOw==, tarball: https://registry.npmjs.org/@tsconfig/node10/-/node10-1.0.11.tgz}
     dev: true
@@ -1558,43 +1546,49 @@ packages:
     resolution: {integrity: sha512-jg+97EGIcY9AGHJJRaaPVgetKDsrTgbRjQ5Msgjh/DQKEFl0DtyRr/VCOyD1T2R1MNeWPK/u7JoGhlDZnKBAfA==, tarball: https://registry.npmjs.org/@types/uuid/-/uuid-9.0.8.tgz}
     dev: true
 
-  /@vitest/expect@1.6.0:
-    resolution: {integrity: sha512-ixEvFVQjycy/oNgHjqsL6AZCDduC+tflRluaHIzKIsdbzkLn2U/iBnVeJwB6HsIjQBdfMR8Z0tRxKUsvFJEeWQ==, tarball: https://registry.npmjs.org/@vitest/expect/-/expect-1.6.0.tgz}
+  /@vitest/expect@2.0.3:
+    resolution: {integrity: sha512-X6AepoOYePM0lDNUPsGXTxgXZAl3EXd0GYe/MZyVE4HzkUqyUVC6S3PrY5mClDJ6/7/7vALLMV3+xD/Ko60Hqg==, tarball: https://registry.npmjs.org/@vitest/expect/-/expect-2.0.3.tgz}
+    dependencies:
+      '@vitest/spy': 2.0.3
+      '@vitest/utils': 2.0.3
+      chai: 5.1.1
+      tinyrainbow: 1.2.0
+    dev: true
+
+  /@vitest/pretty-format@2.0.3:
+    resolution: {integrity: sha512-URM4GLsB2xD37nnTyvf6kfObFafxmycCL8un3OC9gaCs5cti2u+5rJdIflZ2fUJUen4NbvF6jCufwViAFLvz1g==, tarball: https://registry.npmjs.org/@vitest/pretty-format/-/pretty-format-2.0.3.tgz}
     dependencies:
-      '@vitest/spy': 1.6.0
-      '@vitest/utils': 1.6.0
-      chai: 4.4.1
+      tinyrainbow: 1.2.0
     dev: true
 
-  /@vitest/runner@1.6.0:
-    resolution: {integrity: sha512-P4xgwPjwesuBiHisAVz/LSSZtDjOTPYZVmNAnpHHSR6ONrf8eCJOFRvUwdHn30F5M1fxhqtl7QZQUk2dprIXAg==, tarball: https://registry.npmjs.org/@vitest/runner/-/runner-1.6.0.tgz}
+  /@vitest/runner@2.0.3:
+    resolution: {integrity: sha512-EmSP4mcjYhAcuBWwqgpjR3FYVeiA4ROzRunqKltWjBfLNs1tnMLtF+qtgd5ClTwkDP6/DGlKJTNa6WxNK0bNYQ==, tarball: https://registry.npmjs.org/@vitest/runner/-/runner-2.0.3.tgz}
     dependencies:
-      '@vitest/utils': 1.6.0
-      p-limit: 5.0.0
+      '@vitest/utils': 2.0.3
       pathe: 1.1.2
     dev: true
 
-  /@vitest/snapshot@1.6.0:
-    resolution: {integrity: sha512-+Hx43f8Chus+DCmygqqfetcAZrDJwvTj0ymqjQq4CvmpKFSTVteEOBzCusu1x2tt4OJcvBflyHUE0DZSLgEMtQ==, tarball: https://registry.npmjs.org/@vitest/snapshot/-/snapshot-1.6.0.tgz}
+  /@vitest/snapshot@2.0.3:
+    resolution: {integrity: sha512-6OyA6v65Oe3tTzoSuRPcU6kh9m+mPL1vQ2jDlPdn9IQoUxl8rXhBnfICNOC+vwxWY684Vt5UPgtcA2aPFBb6wg==, tarball: https://registry.npmjs.org/@vitest/snapshot/-/snapshot-2.0.3.tgz}
     dependencies:
+      '@vitest/pretty-format': 2.0.3
       magic-string: 0.30.10
       pathe: 1.1.2
-      pretty-format: 29.7.0
     dev: true
 
-  /@vitest/spy@1.6.0:
-    resolution: {integrity: sha512-leUTap6B/cqi/bQkXUu6bQV5TZPx7pmMBKBQiI0rJA8c3pB56ZsaTbREnF7CJfmvAS4V2cXIBAh/3rVwrrCYgw==, tarball: https://registry.npmjs.org/@vitest/spy/-/spy-1.6.0.tgz}
+  /@vitest/spy@2.0.3:
+    resolution: {integrity: sha512-sfqyAw/ypOXlaj4S+w8689qKM1OyPOqnonqOc9T91DsoHbfN5mU7FdifWWv3MtQFf0lEUstEwR9L/q/M390C+A==, tarball: https://registry.npmjs.org/@vitest/spy/-/spy-2.0.3.tgz}
     dependencies:
-      tinyspy: 2.2.1
+      tinyspy: 3.0.0
     dev: true
 
-  /@vitest/utils@1.6.0:
-    resolution: {integrity: sha512-21cPiuGMoMZwiOHa2i4LXkMkMkCGzA+MVFV70jRwHo95dL4x/ts5GZhML1QWuy7yfp3WzK3lRvZi3JnXTYqrBw==, tarball: https://registry.npmjs.org/@vitest/utils/-/utils-1.6.0.tgz}
+  /@vitest/utils@2.0.3:
+    resolution: {integrity: sha512-c/UdELMuHitQbbc/EVctlBaxoYAwQPQdSNwv7z/vHyBKy2edYZaFgptE27BRueZB7eW8po+cllotMNTDpL3HWg==, tarball: https://registry.npmjs.org/@vitest/utils/-/utils-2.0.3.tgz}
     dependencies:
-      diff-sequences: 29.6.3
+      '@vitest/pretty-format': 2.0.3
       estree-walker: 3.0.3
-      loupe: 2.3.7
-      pretty-format: 29.7.0
+      loupe: 3.1.1
+      tinyrainbow: 1.2.0
     dev: true
 
   /@vue/compiler-core@3.4.33:
@@ -1818,8 +1812,9 @@ packages:
     engines: {node: '>=8'}
     dev: true
 
-  /assertion-error@1.1.0:
-    resolution: {integrity: sha512-jgsaNduz+ndvGyFt3uSuWqvy4lCnIJiovtouQN5JZHOKCS2QuhEdbcQHFhVksz2N2U9hXJo8odG7ETyWlEeuDw==, tarball: https://registry.npmjs.org/assertion-error/-/assertion-error-1.1.0.tgz}
+  /assertion-error@2.0.1:
+    resolution: {integrity: sha512-Izi8RQcffqCeNVgFigKli1ssklIbpHnCYc6AknXGYoB6grJqyeby7jv12JUQgmTAnIDnbck1uxksT4dzN3PWBA==, tarball: https://registry.npmjs.org/assertion-error/-/assertion-error-2.0.1.tgz}
+    engines: {node: '>=12'}
     dev: true
 
   /asynckit@0.4.0:
@@ -1903,17 +1898,15 @@ packages:
     engines: {node: '>=10'}
     dev: true
 
-  /chai@4.4.1:
-    resolution: {integrity: sha512-13sOfMv2+DWduEU+/xbun3LScLoqN17nBeTLUsmDfKdoiC1fr0n9PU4guu4AhRcOVFk/sW8LyZWHuhWtQZiF+g==, tarball: https://registry.npmjs.org/chai/-/chai-4.4.1.tgz}
-    engines: {node: '>=4'}
+  /chai@5.1.1:
+    resolution: {integrity: sha512-pT1ZgP8rPNqUgieVaEY+ryQr6Q4HXNg8Ei9UnLUrjN4IA7dvQC5JB+/kxVcPNDHyBcc/26CXPkbNzq3qwrOEKA==, tarball: https://registry.npmjs.org/chai/-/chai-5.1.1.tgz}
+    engines: {node: '>=12'}
     dependencies:
-      assertion-error: 1.1.0
-      check-error: 1.0.3
-      deep-eql: 4.1.4
-      get-func-name: 2.0.2
-      loupe: 2.3.7
-      pathval: 1.1.1
-      type-detect: 4.0.8
+      assertion-error: 2.0.1
+      check-error: 2.1.1
+      deep-eql: 5.0.2
+      loupe: 3.1.1
+      pathval: 2.0.0
     dev: true
 
   /chalk@2.4.2:
@@ -1934,10 +1927,9 @@ packages:
     resolution: {integrity: sha512-mT8iDcrh03qDGRRmoA2hmBJnxpllMR+0/0qlzjqZES6NdiWDcZkCNAk4rPFZ9Q85r27unkiNNg8ZOiwZXBHwcA==, tarball: https://registry.npmjs.org/chardet/-/chardet-0.7.0.tgz}
     dev: true
 
-  /check-error@1.0.3:
-    resolution: {integrity: sha512-iKEoDYaRmd1mxM90a2OEfWhjsjPpYPuQ+lMYsoxB126+t8fw7ySEO48nmDg5COTjxDI65/Y2OWpeEHk3ZOe8zg==, tarball: https://registry.npmjs.org/check-error/-/check-error-1.0.3.tgz}
-    dependencies:
-      get-func-name: 2.0.2
+  /check-error@2.1.1:
+    resolution: {integrity: sha512-OAlb+T7V4Op9OwdkjmguYRqncdlx5JiofwOAUkmTF+jNdHwzTaTs4sRAGpzLF3oOz5xAyDGrPgeIDFQmDOTiJw==, tarball: https://registry.npmjs.org/check-error/-/check-error-2.1.1.tgz}
+    engines: {node: '>= 16'}
     dev: true
 
   /chokidar@3.6.0:
@@ -2012,10 +2004,6 @@ packages:
     engines: {node: '>= 6'}
     dev: true
 
-  /confbox@0.1.7:
-    resolution: {integrity: sha512-uJcB/FKZtBMCJpK8MQji6bJHgu1tixKPxRLeGkNzBoOZzpnZUJm0jm2/sBDWcuBx1dYgxV4JU+g5hmNxCyAmdA==, tarball: https://registry.npmjs.org/confbox/-/confbox-0.1.7.tgz}
-    dev: true
-
   /consola@3.2.3:
     resolution: {integrity: sha512-I5qxpzLv+sJhTVEoLYNcTW+bThDCPsit0vLNKShZx6rLtpilNpmmeTPaeqJb9ZE9dV3DGaeby6Vuhrw38WjeyQ==, tarball: https://registry.npmjs.org/consola/-/consola-3.2.3.tgz}
     engines: {node: ^14.18.0 || >=16.10.0}
@@ -2075,11 +2063,9 @@ packages:
     engines: {node: '>=0.10.0'}
     dev: true
 
-  /deep-eql@4.1.4:
-    resolution: {integrity: sha512-SUwdGfqdKOwxCPeVYjwSyRpJ7Z+fhpwIAtmCUdZIWZ/YP5R9WAsyuSgpLVDi9bjWoN2LXHNss/dk3urXtdQxGg==, tarball: https://registry.npmjs.org/deep-eql/-/deep-eql-4.1.4.tgz}
+  /deep-eql@5.0.2:
+    resolution: {integrity: sha512-h5k/5U50IJJFpzfL6nO9jaaumfjO/f2NjK/oYB2Djzm4p9L+3T9qWpZqZ2hAbLPuuYq9wrU08WQyBTL5GbPk5Q==, tarball: https://registry.npmjs.org/deep-eql/-/deep-eql-5.0.2.tgz}
     engines: {node: '>=6'}
-    dependencies:
-      type-detect: 4.0.8
     dev: true
 
   /delayed-stream@1.0.0:
@@ -2101,11 +2087,6 @@ packages:
     resolution: {integrity: sha512-IayShXAgj/QMXgB0IWmKx+rOPuGMhqm5w6jvFxmVenXKIzRqTAAsbBPT3kWQeGANj3jGgvcvv4yK6SxqYmikgw==, tarball: https://registry.npmjs.org/diff-match-patch/-/diff-match-patch-1.0.5.tgz}
     dev: false
 
-  /diff-sequences@29.6.3:
-    resolution: {integrity: sha512-EjePK1srD3P08o2j4f0ExnylqRs5B9tJjcp9t1krH2qRi8CCdsYfwe9JgSLurFBWwq4uOlipzfk5fHNvwFKr8Q==, tarball: https://registry.npmjs.org/diff-sequences/-/diff-sequences-29.6.3.tgz}
-    engines: {node: ^14.15.0 || ^16.10.0 || >=18.0.0}
-    dev: true
-
   /diff@4.0.2:
     resolution: {integrity: sha512-58lmxKSA4BNyLz+HHMUzlOEpg09FV+ev6ZMe3vJihgdxzgcwZ8VoEEPmALCZG9LmqfVoNMMKpttIYTVG6uDY7A==, tarball: https://registry.npmjs.org/diff/-/diff-4.0.2.tgz}
     engines: {node: '>=0.3.1'}
@@ -2594,10 +2575,6 @@ packages:
     resolution: {integrity: sha512-RdJUflcE3cUzKiMqQgsCu06FPu9UdIJO0beYbPhHN4k6apgJtifcoCtT9bcxOpYBtpD2kCM6Sbzg4CausW/PKQ==, tarball: https://registry.npmjs.org/js-tokens/-/js-tokens-4.0.0.tgz}
     dev: false
 
-  /js-tokens@9.0.0:
-    resolution: {integrity: sha512-WriZw1luRMlmV3LGJaR6QOJjWwgLUTf89OwT2lUOyjX2dJGBwgmIkbcz+7WFZjrZM635JOIR517++e/67CP9dQ==, tarball: https://registry.npmjs.org/js-tokens/-/js-tokens-9.0.0.tgz}
-    dev: true
-
   /js-yaml@3.14.1:
     resolution: {integrity: sha512-okMH7OXXJ7YrN9Ok3/SXrnu4iX9yOk+25nqX4imS2npuvTYDmo/QEZoqwZkYaIDk3jVvBOTOIEgEhaLOynBS9g==, tarball: https://registry.npmjs.org/js-yaml/-/js-yaml-3.14.1.tgz}
     hasBin: true
@@ -2703,14 +2680,6 @@ packages:
       strip-bom: 3.0.0
     dev: true
 
-  /local-pkg@0.5.0:
-    resolution: {integrity: sha512-ok6z3qlYyCDS4ZEU27HaU6x/xZa9Whf8jD4ptH5UZTQYZVYeb9bnZ3ojVhiJNLiXK1Hfc0GNbLXcmZ5plLDDBg==, tarball: https://registry.npmjs.org/local-pkg/-/local-pkg-0.5.0.tgz}
-    engines: {node: '>=14'}
-    dependencies:
-      mlly: 1.7.1
-      pkg-types: 1.1.3
-    dev: true
-
   /locate-character@3.0.0:
     resolution: {integrity: sha512-SW13ws7BjaeJ6p7Q6CO2nchbYEc3X3J6WrmTTDto7yMPqVSZTUyY5Tjbid+Ab8gLnATtygYtiDIJGQRRn2ZOiA==, tarball: https://registry.npmjs.org/locate-character/-/locate-character-3.0.0.tgz}
     dev: false
@@ -2744,8 +2713,8 @@ packages:
       js-tokens: 4.0.0
     dev: false
 
-  /loupe@2.3.7:
-    resolution: {integrity: sha512-zSMINGVYkdpYSOBmLi0D1Uo7JU9nVdQKrHxC8eYlV+9YKK9WePqAlL7lSlorG/U2Fw1w0hTBmaa/jrQ3UbPHtA==, tarball: https://registry.npmjs.org/loupe/-/loupe-2.3.7.tgz}
+  /loupe@3.1.1:
+    resolution: {integrity: sha512-edNu/8D5MKVfGVFRhFf8aAxiTM6Wumfz5XsaatSxlD3w4R1d/WEKUTydCdPGbl9K7QG/Ca3GnDV2sIKIpXRQcw==, tarball: https://registry.npmjs.org/loupe/-/loupe-3.1.1.tgz}
     dependencies:
       get-func-name: 2.0.2
     dev: true
@@ -2856,15 +2825,6 @@ packages:
       num-sort: 2.1.0
     dev: true
 
-  /mlly@1.7.1:
-    resolution: {integrity: sha512-rrVRZRELyQzrIUAVMHxP97kv+G786pHmOKzuFII8zDYahFBS7qnHh2AlYSl1GAHhaMPCz6/oHjVMcfFYgFYHgA==, tarball: https://registry.npmjs.org/mlly/-/mlly-1.7.1.tgz}
-    dependencies:
-      acorn: 8.12.1
-      pathe: 1.1.2
-      pkg-types: 1.1.3
-      ufo: 1.5.4
-    dev: true
-
   /mri@1.2.0:
     resolution: {integrity: sha512-tzzskb3bG8LvYGFF/mDTpq3jpI6Q9wc3LEmBaghu+DdCssd1FakN7Bc0hVNmEyGq1bq3RgfkCb3cmQLpNPOroA==, tarball: https://registry.npmjs.org/mri/-/mri-1.2.0.tgz}
     engines: {node: '>=4'}
@@ -3018,13 +2978,6 @@ packages:
       yocto-queue: 0.1.0
     dev: true
 
-  /p-limit@5.0.0:
-    resolution: {integrity: sha512-/Eaoq+QyLSiXQ4lyYV23f14mZRQcXnxfHrN0vCai+ak9G0pp9iEQukIIZq5NccEvwRB8PUnZT0KsOoDCINS1qQ==, tarball: https://registry.npmjs.org/p-limit/-/p-limit-5.0.0.tgz}
-    engines: {node: '>=18'}
-    dependencies:
-      yocto-queue: 1.1.1
-    dev: true
-
   /p-locate@4.1.0:
     resolution: {integrity: sha512-R79ZZ/0wAxKGu3oYMlz8jy/kbhsNrS7SKZ7PxEHBgJ5+F2mtFW2fK2cOtBh1cHYkQsbzFV7I+EoRKe6Yt0oK7A==, tarball: https://registry.npmjs.org/p-locate/-/p-locate-4.1.0.tgz}
     engines: {node: '>=8'}
@@ -3108,8 +3061,9 @@ packages:
     resolution: {integrity: sha512-whLdWMYL2TwI08hn8/ZqAbrVemu0LNaNNJZX73O6qaIdCTfXutsLhMkjdENX0qhsQ9uIimo4/aQOmXkoon2nDQ==, tarball: https://registry.npmjs.org/pathe/-/pathe-1.1.2.tgz}
     dev: true
 
-  /pathval@1.1.1:
-    resolution: {integrity: sha512-Dp6zGqpTdETdR63lehJYPeIOqpiNBNtc7BpWSLrOje7UaIsE5aY92r/AunQA7rsXvet3lrJ3JnZX29UPTKXyKQ==, tarball: https://registry.npmjs.org/pathval/-/pathval-1.1.1.tgz}
+  /pathval@2.0.0:
+    resolution: {integrity: sha512-vE7JKRyES09KiunauX7nd2Q9/L7lhok4smP9RZTDeD4MVs72Dp2qNFVz39Nz5a0FVEW0BJR6C0DYrq6unoziZA==, tarball: https://registry.npmjs.org/pathval/-/pathval-2.0.0.tgz}
+    engines: {node: '>= 14.16'}
     dev: true
 
   /periscopic@3.1.0:
@@ -3145,14 +3099,6 @@ packages:
       find-up: 4.1.0
     dev: true
 
-  /pkg-types@1.1.3:
-    resolution: {integrity: sha512-+JrgthZG6m3ckicaOB74TwQ+tBWsFl3qVQg7mN8ulwSOElJ7gBhKzj2VkCPnZ4NlF6kEquYU+RIYNVAvzd54UA==, tarball: https://registry.npmjs.org/pkg-types/-/pkg-types-1.1.3.tgz}
-    dependencies:
-      confbox: 0.1.7
-      mlly: 1.7.1
-      pathe: 1.1.2
-    dev: true
-
   /postcss-load-config@6.0.1:
     resolution: {integrity: sha512-oPtTM4oerL+UXmx+93ytZVN82RrlY/wPUV8IeDxFrzIjXOLF1pN+EmKPLbubvKHT2HC20xXsCAH2Z+CKV6Oz/g==, tarball: https://registry.npmjs.org/postcss-load-config/-/postcss-load-config-6.0.1.tgz}
     engines: {node: '>= 18'}
@@ -3198,15 +3144,6 @@ packages:
     hasBin: true
     dev: true
 
-  /pretty-format@29.7.0:
-    resolution: {integrity: sha512-Pdlw/oPxN+aXdmM9R00JVC9WVFoCLTKJvDVLgmJ+qAffBMxsV85l/Lu7sNx4zSzPyoL2euImuEwHhOXdEgNFZQ==, tarball: https://registry.npmjs.org/pretty-format/-/pretty-format-29.7.0.tgz}
-    engines: {node: ^14.15.0 || ^16.10.0 || >=18.0.0}
-    dependencies:
-      '@jest/schemas': 29.6.3
-      ansi-styles: 5.2.0
-      react-is: 18.3.1
-    dev: true
-
   /proxy-from-env@1.1.0:
     resolution: {integrity: sha512-D+zkORCbA9f1tdWRK0RaCR3GPv50cMxcrz4X8k5LTSUD1Dkw47mKJEZQNunItRTkWwgtaUSo1RVFRIG9ZXiFYg==, tarball: https://registry.npmjs.org/proxy-from-env/-/proxy-from-env-1.1.0.tgz}
     dev: true
@@ -3224,10 +3161,6 @@ packages:
     resolution: {integrity: sha512-NuaNSa6flKT5JaSYQzJok04JzTL1CA6aGhv5rfLW3PgqA+M2ChpZQnAC8h8i4ZFkBS8X5RqkDBHA7r4hej3K9A==, tarball: https://registry.npmjs.org/queue-microtask/-/queue-microtask-1.2.3.tgz}
     dev: true
 
-  /react-is@18.3.1:
-    resolution: {integrity: sha512-/LLMVyas0ljjAtoYiPqYiL8VWXzUUdThrmU5+n20DZv+a+ClRoevUzw5JxU+Ieh5/c87ytoTBV9G1FiKfNJdmg==, tarball: https://registry.npmjs.org/react-is/-/react-is-18.3.1.tgz}
-    dev: true
-
   /react@18.3.1:
     resolution: {integrity: sha512-wS+hAgJShR0KhEvPJArfuPVN1+Hz1t0Y6n5jLrGQbkb4urgPE/0Rve+1kMB1v/oWgHgm4WIcV+i7F2pTVj+2iQ==, tarball: https://registry.npmjs.org/react/-/react-18.3.1.tgz}
     engines: {node: '>=0.10.0'}
@@ -3445,12 +3378,6 @@ packages:
     engines: {node: '>=12'}
     dev: true
 
-  /strip-literal@2.1.0:
-    resolution: {integrity: sha512-Op+UycaUt/8FbN/Z2TWPBLge3jWrP3xj10f3fnYxf052bKuS3EKs1ZQcVGjnEMdsNVAM+plXRdmjrZ/KgG3Skw==, tarball: https://registry.npmjs.org/strip-literal/-/strip-literal-2.1.0.tgz}
-    dependencies:
-      js-tokens: 9.0.0
-    dev: true
-
   /sucrase@3.35.0:
     resolution: {integrity: sha512-8EbVDiu9iN/nESwxeSxDKe0dunta1GOlHufmSSXxMD2z2/tMZpDMpvXQGsc+ajGo8y2uYUmixaSRUc/QPoQ0GA==, tarball: https://registry.npmjs.org/sucrase/-/sucrase-3.35.0.tgz}
     engines: {node: '>=16 || 14 >=14.17'}
@@ -3536,13 +3463,18 @@ packages:
     resolution: {integrity: sha512-1/eK7zUnIklz4JUUlL+658n58XO2hHLQfSk1Zf2LKieUjxidN16eKFEoDEfjHc3ohofSSqK3X5yO6VGb6iW8Lw==, tarball: https://registry.npmjs.org/tinybench/-/tinybench-2.8.0.tgz}
     dev: true
 
-  /tinypool@0.8.4:
-    resolution: {integrity: sha512-i11VH5gS6IFeLY3gMBQ00/MmLncVP7JLXOw1vlgkytLmJK7QnEr7NXf0LBdxfmNPAeyetukOk0bOYrJrFGjYJQ==, tarball: https://registry.npmjs.org/tinypool/-/tinypool-0.8.4.tgz}
+  /tinypool@1.0.0:
+    resolution: {integrity: sha512-KIKExllK7jp3uvrNtvRBYBWBOAXSX8ZvoaD8T+7KB/QHIuoJW3Pmr60zucywjAlMb5TeXUkcs/MWeWLu0qvuAQ==, tarball: https://registry.npmjs.org/tinypool/-/tinypool-1.0.0.tgz}
+    engines: {node: ^18.0.0 || >=20.0.0}
+    dev: true
+
+  /tinyrainbow@1.2.0:
+    resolution: {integrity: sha512-weEDEq7Z5eTHPDh4xjX789+fHfF+P8boiFB+0vbWzpbnbsEr/GRaohi/uMKxg8RZMXnl1ItAi/IUHWMsjDV7kQ==, tarball: https://registry.npmjs.org/tinyrainbow/-/tinyrainbow-1.2.0.tgz}
     engines: {node: '>=14.0.0'}
     dev: true
 
-  /tinyspy@2.2.1:
-    resolution: {integrity: sha512-KYad6Vy5VDWV4GH3fjpseMQ/XU2BhIYP7Vzd0LG44qRWm/Yt2WCOTicFdvmgo6gWaqooMQCawTtILVQJupKu7A==, tarball: https://registry.npmjs.org/tinyspy/-/tinyspy-2.2.1.tgz}
+  /tinyspy@3.0.0:
+    resolution: {integrity: sha512-q5nmENpTHgiPVd1cJDDc9cVoYN5x4vCvwT3FMilvKPKneCBZAxn2YWQjDF0UMcE9k0Cay1gBiDfTMU0g+mPMQA==, tarball: https://registry.npmjs.org/tinyspy/-/tinyspy-3.0.0.tgz}
     engines: {node: '>=14.0.0'}
     dev: true
 
@@ -3661,20 +3593,11 @@ packages:
       - yaml
     dev: true
 
-  /type-detect@4.0.8:
-    resolution: {integrity: sha512-0fr/mIH1dlO+x7TlcMy+bIDqKPsw/70tVyeHW787goQjhmqaZe10uwLujubK9q9Lg6Fiho1KUKDYz0Z7k7g5/g==, tarball: https://registry.npmjs.org/type-detect/-/type-detect-4.0.8.tgz}
-    engines: {node: '>=4'}
-    dev: true
-
   /typescript@5.5.3:
     resolution: {integrity: sha512-/hreyEujaB0w76zKo6717l3L0o/qEUtRgdvUBvlkhoWeOVMjMuHNHk0BRBzikzuGDqNmPQbg5ifMEqsHLiIUcQ==, tarball: https://registry.npmjs.org/typescript/-/typescript-5.5.3.tgz}
     engines: {node: '>=14.17'}
     hasBin: true
 
-  /ufo@1.5.4:
-    resolution: {integrity: sha512-UsUk3byDzKd04EyoZ7U4DOlxQaD14JUKQl6/P7wiX4FNvUfm3XL246n9W5AmqwW5RSFJ27NAuM0iLscAOYUiGQ==, tarball: https://registry.npmjs.org/ufo/-/ufo-1.5.4.tgz}
-    dev: true
-
   /undici-types@5.26.5:
     resolution: {integrity: sha512-JlCMO+ehdEIKqlFxk6IfVoAUVmgz7cU7zD/h9XZ0qzeosSHmUJVOzSQvvYSYWXkFXC+IfLKSIffhv0sVZup6pA==, tarball: https://registry.npmjs.org/undici-types/-/undici-types-5.26.5.tgz}
     dev: true
@@ -3701,15 +3624,15 @@ packages:
     resolution: {integrity: sha512-wa7YjyUGfNZngI/vtK0UHAN+lgDCxBPCylVXGp0zu59Fz5aiGtNXaq3DhIov063MorB+VfufLh3JlF2KdTK3xg==, tarball: https://registry.npmjs.org/v8-compile-cache-lib/-/v8-compile-cache-lib-3.0.1.tgz}
     dev: true
 
-  /vite-node@1.6.0(@types/node@20.14.11):
-    resolution: {integrity: sha512-de6HJgzC+TFzOu0NTC4RAIsyf/DY/ibWDYQUcuEA84EMHhcefTUGkjFHKKEJhQN4A+6I0u++kr3l36ZF2d7XRw==, tarball: https://registry.npmjs.org/vite-node/-/vite-node-1.6.0.tgz}
+  /vite-node@2.0.3(@types/node@20.14.11):
+    resolution: {integrity: sha512-14jzwMx7XTcMB+9BhGQyoEAmSl0eOr3nrnn+Z12WNERtOvLN+d2scbRUvyni05rT3997Bg+rZb47NyP4IQPKXg==, tarball: https://registry.npmjs.org/vite-node/-/vite-node-2.0.3.tgz}
     engines: {node: ^18.0.0 || >=20.0.0}
     hasBin: true
     dependencies:
       cac: 6.7.14
       debug: 4.3.5
       pathe: 1.1.2
-      picocolors: 1.0.1
+      tinyrainbow: 1.2.0
       vite: 5.3.4(@types/node@20.14.11)
     transitivePeerDependencies:
       - '@types/node'
@@ -3758,15 +3681,15 @@ packages:
       fsevents: 2.3.3
     dev: true
 
-  /vitest@1.6.0(@types/node@20.14.11):
-    resolution: {integrity: sha512-H5r/dN06swuFnzNFhq/dnz37bPXnq8xB2xB5JOVk8K09rUtoeNN+LHWkoQ0A/i3hvbUKKcCei9KpbxqHMLhLLA==, tarball: https://registry.npmjs.org/vitest/-/vitest-1.6.0.tgz}
+  /vitest@2.0.3(@types/node@20.14.11):
+    resolution: {integrity: sha512-o3HRvU93q6qZK4rI2JrhKyZMMuxg/JRt30E6qeQs6ueaiz5hr1cPj+Sk2kATgQzMMqsa2DiNI0TIK++1ULx8Jw==, tarball: https://registry.npmjs.org/vitest/-/vitest-2.0.3.tgz}
     engines: {node: ^18.0.0 || >=20.0.0}
     hasBin: true
     peerDependencies:
       '@edge-runtime/vm': '*'
       '@types/node': ^18.0.0 || >=20.0.0
-      '@vitest/browser': 1.6.0
-      '@vitest/ui': 1.6.0
+      '@vitest/browser': 2.0.3
+      '@vitest/ui': 2.0.3
       happy-dom: '*'
       jsdom: '*'
     peerDependenciesMeta:
@@ -3783,26 +3706,25 @@ packages:
       jsdom:
         optional: true
     dependencies:
+      '@ampproject/remapping': 2.3.0
       '@types/node': 20.14.11
-      '@vitest/expect': 1.6.0
-      '@vitest/runner': 1.6.0
-      '@vitest/snapshot': 1.6.0
-      '@vitest/spy': 1.6.0
-      '@vitest/utils': 1.6.0
-      acorn-walk: 8.3.3
-      chai: 4.4.1
+      '@vitest/expect': 2.0.3
+      '@vitest/pretty-format': 2.0.3
+      '@vitest/runner': 2.0.3
+      '@vitest/snapshot': 2.0.3
+      '@vitest/spy': 2.0.3
+      '@vitest/utils': 2.0.3
+      chai: 5.1.1
       debug: 4.3.5
       execa: 8.0.1
-      local-pkg: 0.5.0
       magic-string: 0.30.10
       pathe: 1.1.2
-      picocolors: 1.0.1
       std-env: 3.7.0
-      strip-literal: 2.1.0
       tinybench: 2.8.0
-      tinypool: 0.8.4
+      tinypool: 1.0.0
+      tinyrainbow: 1.2.0
       vite: 5.3.4(@types/node@20.14.11)
-      vite-node: 1.6.0(@types/node@20.14.11)
+      vite-node: 2.0.3(@types/node@20.14.11)
       why-is-node-running: 2.3.0
     transitivePeerDependencies:
       - less
@@ -3941,11 +3863,6 @@ packages:
     engines: {node: '>=10'}
     dev: true
 
-  /yocto-queue@1.1.1:
-    resolution: {integrity: sha512-b4JR1PFR10y1mKjhHY9LaGo6tmrgjit7hxVIeAmyMw3jegXR4dhYqLaQF5zMXZxY7tLpMyJeLjr1C4rLmkVe8g==, tarball: https://registry.npmjs.org/yocto-queue/-/yocto-queue-1.1.1.tgz}
-    engines: {node: '>=12.20'}
-    dev: true
-
   /zod-to-json-schema@3.22.5(zod@3.23.8):
     resolution: {integrity: sha512-+akaPo6a0zpVCCseDed504KBJUQpEW5QZw7RMneNmKw+fGaML1Z9tUNLnHHAC8x6dzVRO1eB2oEMyZRnuBZg7Q==, tarball: https://registry.npmjs.org/zod-to-json-schema/-/zod-to-json-schema-3.22.5.tgz}
     peerDependencies:
diff --git a/src/agent.test.ts b/src/agent.test.ts
index 5ad35f7..f943f19 100644
--- a/src/agent.test.ts
+++ b/src/agent.test.ts
@@ -1,5 +1,10 @@
 import { test, expect, vi } from 'vitest';
-import { createAgent, type AIAdapter } from './';
+import {
+  AgentGenerateTextResult,
+  AgentMessage,
+  createAgent,
+  type AIAdapter,
+} from './';
 import { createActor, createMachine } from 'xstate';
 import { GenerateTextResult } from 'ai';
 import { z } from 'zod';
@@ -297,7 +302,7 @@ test('You can listen for plan events', async () => {
               },
             },
           ],
-        } as any as GenerateTextResult<any>;
+        } as any as AgentGenerateTextResult;
       },
       streamText: {} as any,
     },
@@ -357,3 +362,145 @@ test('agent.types provides context and event types', () => {
   // @ts-expect-error
   agent.types.context satisfies { score: string };
 });
+
+test.each(['generateText', 'streamText'] as const)(
+  'can provide a correlation ID (%s)',
+  async (method) => {
+    const agent = createAgent({
+      model: {} as any,
+      events: {},
+      adapter: {
+        [method]: async (opts) => {
+          const res = {
+            text: 'response',
+          };
+
+          opts.onFinish?.(res);
+
+          return res as AgentGenerateTextResult;
+        },
+      } as any as AIAdapter,
+    });
+
+    const promise = new Promise<AgentMessage>((res) => {
+      agent.onMessage((msg) => {
+        if (msg.role === 'assistant') {
+          res(msg);
+        }
+      });
+    });
+
+    await agent[method]({
+      prompt: 'hi',
+      correlationId: 'c-1',
+    });
+
+    const msg = await promise;
+
+    expect(msg.correlationId).toBe('c-1');
+    expect(msg.parentCorrelationId).toBe(undefined);
+  }
+);
+
+test.each(['generateText', 'streamText'] as const)(
+  'correlation IDs are automatically generated if not provided (%s)',
+  async (method) => {
+    const agent = createAgent({
+      model: {} as any,
+      events: {},
+      adapter: {
+        [method]: async (opts) => {
+          const res = {
+            text: 'response',
+          };
+
+          opts.onFinish?.(res);
+
+          return res as AgentGenerateTextResult;
+        },
+      } as any as AIAdapter,
+    });
+
+    await agent[method]({
+      prompt: 'hi',
+    });
+
+    const messages = agent.getMessages();
+
+    expect(messages[0]?.correlationId).toEqual(expect.stringMatching(/.+/));
+    expect(messages[0]?.role).toBe('user');
+    expect(messages[1]?.correlationId).toEqual(expect.stringMatching(/.+/));
+    expect(messages[1]?.role).toBe('assistant');
+
+    expect(messages[0]!.correlationId).toEqual(messages[1]!.correlationId);
+  }
+);
+
+test.each(['generateText', 'streamText'] as const)(
+  'can provide a parent correlation ID (%s)',
+  async (method) => {
+    const agent = createAgent({
+      model: {} as any,
+      events: {},
+      adapter: {
+        [method]: async (opts) => {
+          const res = {
+            text: 'response',
+          };
+
+          opts.onFinish?.(res);
+
+          return res as AgentGenerateTextResult;
+        },
+      } as any as AIAdapter,
+    });
+
+    await agent[method]({
+      prompt: 'hi',
+      correlationId: 'c-1',
+      parentCorrelationId: 'c-0',
+    });
+
+    const msg = agent.getMessages().find((msg) => msg.role === 'assistant')!;
+
+    expect(msg.correlationId).toBe('c-1');
+    expect(msg.parentCorrelationId).toBe('c-0');
+  }
+);
+
+test.each(['generateText', 'streamText'] as const)(
+  'can add feedback to a correlation (%s)',
+  async (method) => {
+    const agent = createAgent({
+      name: 'test',
+      model: {} as any,
+      events: {},
+      adapter: {
+        [method]: async (opts) => {
+          const res = {
+            text: 'response',
+          };
+
+          opts.onFinish?.(res);
+
+          return res as AgentGenerateTextResult;
+        },
+      } as any as AIAdapter,
+    });
+
+    const res = await agent[method]({
+      prompt: 'test',
+    });
+
+    agent.addFeedback({
+      correlationId: res.correlationId,
+      reward: -1,
+    });
+
+    const message = agent.getMessages()[0]!;
+    const feedback = agent.getFeedback()[0]!;
+
+    expect(message.correlationId).toBeDefined();
+    expect(feedback.correlationId).toEqual(message.correlationId);
+  }
+);
diff --git a/src/agent.ts b/src/agent.ts
index 3fb0b03..3ff0e96 100644
--- a/src/agent.ts
+++ b/src/agent.ts
@@ -22,6 +22,7 @@ import {
   AgentMemoryContext,
   AgentObservation,
   ContextFromZodContextMapping,
+  AgentFeedback,
 } from './types';
 import { simplePlanner } from './planners/simplePlanner';
 import { agentGenerateText, agentStreamText } from './text';
@@ -73,12 +74,13 @@ export const agentLogic: AgentLogic<AnyEventObject> = fromTransition(
     }
     return state;
   },
-  {
-    feedback: [],
-    messages: [],
-    observations: [],
-    plans: [],
-  } as AgentMemoryContext
+  () =>
+    ({
+      feedback: [],
+      messages: [],
+      observations: [],
+      plans: [],
+    } as AgentMemoryContext)
 );
 
 export function createAgent<
@@ -140,8 +142,6 @@ export function createAgent<
   logic?: AgentLogic<TEvents>;
   adapter?: AIAdapter;
 } & GenerateTextOptions): Agent<TContext, TEvents> {
-  const messageHistoryListeners: Observer<AgentMessage>[] = [];
-
   const agent = createActor(logic) as unknown as Agent<TContext, TEvents>;
   agent.events = events;
   agent.model = model;
@@ -155,7 +155,7 @@ export function createAgent<
   agent.memory = getMemory ? getMemory(agent) : undefined;
 
   agent.onMessage = (callback) => {
-    messageHistoryListeners.push(toObserver(callback));
+    agent.on('message', (ev) => callback(ev.message));
   };
 
   agent.decide = (opts) => {
@@ -168,7 +168,8 @@ export function createAgent<
       id: messageInput.id ?? randomId(),
       timestamp: messageInput.timestamp ?? Date.now(),
       sessionId: agent.sessionId,
-    };
+      correlationId: messageInput.correlationId ?? randomId(),
+    } satisfies AgentMessage;
     agent.send({
       type: 'agent.message',
       message,
@@ -185,9 +186,11 @@ export function createAgent<
   agent.addFeedback = (feedbackInput) => {
     const feedback = {
       ...feedbackInput,
+      attributes: { ...feedbackInput.attributes },
+      reward: feedbackInput.reward ?? 0,
       timestamp: feedbackInput.timestamp ?? Date.now(),
       sessionId: agent.sessionId,
-    };
+    } satisfies AgentFeedback;
     agent.send({
       type: 'agent.feedback',
       feedback,
diff --git a/src/planners/simplePlanner.ts b/src/planners/simplePlanner.ts
index 10a36d3..606d3a9 100644
--- a/src/planners/simplePlanner.ts
+++ b/src/planners/simplePlanner.ts
@@ -1,4 +1,4 @@
-import { CoreTool, tool } from 'ai';
+import { type CoreTool, tool } from 'ai';
 import {
   AgentPlan,
   AgentPlanInput,
@@ -118,7 +118,6 @@ export async function simplePlanner<T extends AnyAgent>(
   const singleResult = result.toolResults[0];
 
   if (!singleResult) {
-    console.log(toolMap);
     // TODO: retries?
     console.warn('No tool call results returned');
     return undefined;
diff --git a/src/text.ts b/src/text.ts
index 84d573e..70e7fd0 100644
--- a/src/text.ts
+++ b/src/text.ts
@@ -1,12 +1,9 @@
-import type {
-  CoreMessage,
-  CoreTool,
-  GenerateTextResult,
-  StreamTextResult,
-} from 'ai';
+import type { CoreMessage, CoreTool, GenerateTextResult } from 'ai';
 import {
   AgentGenerateTextOptions,
+  AgentGenerateTextResult,
   AgentStreamTextOptions,
+  AgentStreamTextResult,
   AnyAgent,
 } from './types';
 import { defaultTextTemplate } from './templates/defaultText';
@@ -51,11 +48,13 @@ export async function getMessages(
 export async function agentGenerateText<T extends AnyAgent>(
   agent: T,
   options: AgentGenerateTextOptions
-) {
+): Promise<AgentGenerateTextResult> {
   const resolvedOptions = {
     ...agent.defaultOptions,
     ...options,
+    correlationId: options.correlationId ?? randomId(),
   };
+  // Generate a correlation ID if one is not provided
   const template = resolvedOptions.template ?? defaultTextTemplate;
   // TODO: check if messages was provided instead
   const id = randomId();
@@ -76,6 +75,8 @@ export async function agentGenerateText<T extends AnyAgent>(
     role: 'user',
     content: promptWithContext,
     timestamp: Date.now(),
+    correlationId: resolvedOptions.correlationId,
+    parentCorrelationId: resolvedOptions.parentCorrelationId,
   });
 
   const result = await agent.adapter.generateText({
@@ -91,18 +92,25 @@ export async function agentGenerateText<T extends AnyAgent>(
     timestamp: Date.now(),
     responseId: id,
     result,
+    correlationId: resolvedOptions.correlationId,
+    parentCorrelationId: resolvedOptions.parentCorrelationId,
   });
 
-  return result;
+  return {
+    ...result,
+    parentCorrelationId: resolvedOptions.parentCorrelationId,
+    correlationId: resolvedOptions.correlationId,
+  };
 }
 
 export async function agentStreamText(
   agent: AnyAgent,
   options: AgentStreamTextOptions
-): Promise<StreamTextResult<any>> {
+): Promise<AgentStreamTextResult> {
   const resolvedOptions = {
     ...agent.defaultOptions,
     ...options,
+    correlationId: options.correlationId ?? randomId(),
   };
   const template = resolvedOptions.template ?? defaultTextTemplate;
 
@@ -124,6 +132,8 @@ export async function agentStreamText(
     content: promptWithContext,
     id,
     timestamp: Date.now(),
+    correlationId: resolvedOptions.correlationId,
+    parentCorrelationId: resolvedOptions.parentCorrelationId,
   });
 
   const result = await agent.adapter.streamText({
@@ -149,11 +159,17 @@ export async function agentStreamText(
         id: randomId(),
         timestamp: Date.now(),
         responseId: id,
+        correlationId: resolvedOptions.correlationId,
+        parentCorrelationId: resolvedOptions.parentCorrelationId,
       });
     },
   });
 
-  return result;
+  return {
+    ...result,
+    parentCorrelationId: resolvedOptions.parentCorrelationId,
+    correlationId: resolvedOptions.correlationId,
+  } as unknown as AgentStreamTextResult; // TODO: fix
 }
 
 export function fromTextStream<T extends AnyAgent>(
diff --git a/src/types.ts b/src/types.ts
index 7fc9ce5..4636c24 100644
--- a/src/types.ts
+++ b/src/types.ts
@@ -128,18 +128,25 @@ export type AgentDecideOptions = {
 >;
 
 export interface AgentFeedback {
-  goal: string;
-  observationId: string;
+  goal?: string;
+  observationId?: string;
+  /**
+   * The message correlation that the feedback is relevant for
+   */
+  correlationId?: string;
   attributes: Record<string, any>;
+  reward: number;
   timestamp: number;
   sessionId: string;
 }
 
 export interface AgentFeedbackInput {
-  goal: string;
-  observationId: string; // Observation ID;
-  attributes: Record<string, any>;
+  goal?: string;
+  observationId?: string;
+  correlationId?: string;
+  attributes?: Record<string, any>;
   timestamp?: number;
+  reward?: number;
 }
 
 export type AgentMessage = CoreMessage & {
@@ -152,6 +159,8 @@ export type AgentMessage = CoreMessage & {
   responseId?: string;
   result?: GenerateTextResult<any>;
   sessionId: string;
+  correlationId: string;
+  parentCorrelationId?: string;
 };
 
 export type AgentMessageInput = CoreMessage & {
@@ -162,6 +171,8 @@ export type AgentMessageInput = CoreMessage & {
    * which message this message is responding to, if any.
    */
   responseId?: string;
+  correlationId?: string;
+  parentCorrelationId?: string;
   result?: GenerateTextResult<any>;
 };
 
@@ -293,12 +304,12 @@ export type Agent<TContext, TEvents extends EventObject> = ActorRefFrom<
   // Generate text
   generateText: (
     options: AgentGenerateTextOptions
-  ) => Promise<GenerateTextResult<Record<string, any>>>;
+  ) => Promise<AgentGenerateTextResult>;
 
   // Stream text
   streamText: (
     options: AgentStreamTextOptions
-  ) => Promise<StreamTextResult<Record<string, CoreTool<any, any>>>>;
+  ) => Promise<AgentStreamTextResult>;
 
   addObservation: (
     observationInput: AgentObservationInput
@@ -395,15 +406,22 @@ export type Agent<TContext, TEvents extends EventObject> = ActorRefFrom<
 
 export type AnyAgent = Agent<any, any>;
 
-export type FromAgent<T> = T | ((self: AnyAgent) => T | Promise<T>);
+export type FromAgent<T> = T | ((agent: AnyAgent) => T | Promise<T>);
 
-export interface CommonTextOptions {
+export type CommonTextOptions = {
   prompt: FromAgent<string>;
   model?: LanguageModel;
   context?: Record<string, any>;
   messages?: FromAgent<CoreMessage[]>;
   template?: PromptTemplate<any>;
-}
+  correlationId?: string;
+  parentCorrelationId?: string;
+};
+
+export type TextResultMeta = {
+  correlationId: string;
+  parentCorrelationId?: string;
+};
 
 export type AgentGenerateTextOptions = Omit<
   GenerateTextOptions,
@@ -411,12 +429,16 @@ export type AgentGenerateTextOptions = Omit<
 > &
   CommonTextOptions;
 
+export type AgentGenerateTextResult = GenerateTextResult<any> & TextResultMeta;
+
 export type AgentStreamTextOptions = Omit<
   StreamTextOptions,
   'model' | 'prompt' | 'messages'
 > &
   CommonTextOptions;
 
+export type AgentStreamTextResult = StreamTextResult<any> & TextResultMeta;
+
 export interface ObservedState {
   /**
    * The current state value of the state machine, e.g.