diff --git a/CHANGELOG.md b/CHANGELOG.md
index de32c743..6783372b 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -10,6 +10,7 @@
 - Add support for new model interface [#229](https://github.com/gohypermode/runtime/pull/229)
 - Add sequential vector search [#240](https://github.com/hypermodeAI/runtime/pull/240)
 - Update Hypermode-hosted model endpoint URL [#242](https://github.com/gohypermode/runtime/pull/242)
+- Fix bug caused by #226 [#243](https://github.com/gohypermode/runtime/pull/243)
 
 ## 2024-06-03 - Version 0.8.2
 
diff --git a/graphql/datasource/source.go b/graphql/datasource/source.go
index f07958c7..a17dcbd0 100644
--- a/graphql/datasource/source.go
+++ b/graphql/datasource/source.go
@@ -26,11 +26,6 @@ type callInfo struct {
 	Parameters map[string]any `json:"data"`
 }
 
-type FunctionOutput struct {
-	ExecutionId string
-	Buffers     utils.OutputBuffers
-}
-
 type Source struct{}
 
 func (s Source) Load(ctx context.Context, input []byte, writer io.Writer) error {
@@ -57,19 +52,14 @@ func (s Source) Load(ctx context.Context, input []byte, writer io.Writer) error
 func (s Source) callFunction(ctx context.Context, callInfo callInfo) (any, []resolve.GraphQLError, error) {
 	// Call the function
 	info, err := wasmhost.CallFunctionWithParametersMap(ctx, callInfo.Function.Name, callInfo.Parameters)
-	if err != nil {
-		return nil, nil, err
-	}
+	// NOTE: don't return the error here, as we want to capture function errors in the response.
 
-	// Store the Execution ID and output buffers in the context
-	outputMap := ctx.Value(utils.FunctionOutputContextKey).(map[string]FunctionOutput)
-	outputMap[callInfo.Function.AliasOrName()] = FunctionOutput{
-		ExecutionId: info.ExecutionId,
-		Buffers:     *info.Buffers,
-	}
+	// Store the execution info into the function output map.
+	outputMap := ctx.Value(utils.FunctionOutputContextKey).(map[string]*wasmhost.ExecutionInfo)
+	outputMap[callInfo.Function.AliasOrName()] = info
 
-	// Transform messages (and error lines in the output buffers) to GraphQL errors
-	messages := append(info.Messages, utils.TransformConsoleOutput(*info.Buffers)...)
+	// Transform messages (and error lines in the output buffers) to GraphQL errors.
+	messages := append(info.Messages, utils.TransformConsoleOutput(info.Buffers)...)
 	gqlErrors := transformErrors(messages, callInfo)
 
 	return info.Result, gqlErrors, err
diff --git a/graphql/graphql.go b/graphql/graphql.go
index e7af22cc..50a2157a 100644
--- a/graphql/graphql.go
+++ b/graphql/graphql.go
@@ -9,11 +9,11 @@ import (
 	"fmt"
 	"net/http"
 
-	"hmruntime/graphql/datasource"
 	"hmruntime/graphql/engine"
 	"hmruntime/logger"
 	"hmruntime/pluginmanager"
 	"hmruntime/utils"
+	"hmruntime/wasmhost"
 
 	"github.com/buger/jsonparser"
 	eng "github.com/wundergraph/graphql-go-tools/execution/engine"
@@ -57,7 +57,7 @@ func HandleGraphQLRequest(w http.ResponseWriter, r *http.Request) {
 	}
 
 	// Create the output map
-	output := map[string]datasource.FunctionOutput{}
+	output := map[string]*wasmhost.ExecutionInfo{}
 	ctx = context.WithValue(ctx, utils.FunctionOutputContextKey, output)
 
 	// Set tracing options
@@ -111,7 +111,7 @@ func HandleGraphQLRequest(w http.ResponseWriter, r *http.Request) {
 	_, _ = w.Write(response)
 }
 
-func addOutputToResponse(response []byte, output map[string]datasource.FunctionOutput) ([]byte, error) {
+func addOutputToResponse(response []byte, output map[string]*wasmhost.ExecutionInfo) ([]byte, error) {
 
 	type invocationInfo struct {
 		ExecutionId string             `json:"executionId"`
diff --git a/utils/console.go b/utils/console.go
index 7b347b55..31d6d68d 100644
--- a/utils/console.go
+++ b/utils/console.go
@@ -18,7 +18,7 @@ func (l LogMessage) IsError() bool {
 	return l.Level == "error" || l.Level == "fatal"
 }
 
-func TransformConsoleOutput(buffers OutputBuffers) []LogMessage {
+func TransformConsoleOutput(buffers *OutputBuffers) []LogMessage {
 	return append(transformConsoleOutputLines(&buffers.StdOut), transformConsoleOutputLines(&buffers.StdErr)...)
 }