diff --git a/src/main/java/de/extio/lmlib/agent/BaseAgent.java b/src/main/java/de/extio/lmlib/agent/BaseAgent.java
index 86e0fd7..066a006 100644
--- a/src/main/java/de/extio/lmlib/agent/BaseAgent.java
+++ b/src/main/java/de/extio/lmlib/agent/BaseAgent.java
@@ -142,6 +142,9 @@ public interface BaseAgent {
 					}
 					
 					this.postProcess(split.context());
+					if (split.context().getAgentContextUpdateConsumer() != null) {
+						split.context().getAgentContextUpdateConsumer().accept(split.context());
+					}
 					
 					finishedSplits.add(split);
 				}
diff --git a/src/main/java/de/extio/lmlib/agent/responsehandler/TextAgentResponseHandler.java b/src/main/java/de/extio/lmlib/agent/responsehandler/TextAgentResponseHandler.java
index 18c108b..f878132 100755
--- a/src/main/java/de/extio/lmlib/agent/responsehandler/TextAgentResponseHandler.java
+++ b/src/main/java/de/extio/lmlib/agent/responsehandler/TextAgentResponseHandler.java
@@ -22,18 +22,21 @@ public class TextAgentResponseHandler implements StreamedAgentResponseHandler {
 	
 	private final Function<String, String> transformer;
 	
+	private final Consumer<AgentContext> beforeStream;
+
 	private final Function<String, String> chunkTransformer;
 
 	private final Consumer<AgentContext> afterChunkUpdate;
 	
 	public TextAgentResponseHandler(final String key) {
-		this(key, null, null, null);
+		this(key, null, null, null, null);
 	}
 	
-	public TextAgentResponseHandler(final String key, final Function<String, String> transformer, final Function<String, String> chunkTransformer, final Consumer<AgentContext> afterChunkUpdate) {
+	public TextAgentResponseHandler(final String key, final Function<String, String> transformer, final Consumer<AgentContext> beforeStream, final Function<String, String> chunkTransformer, final Consumer<AgentContext> afterChunkUpdate) {
 		this.key = key;
 		this.chunkKey = key.concat("_chunk");
 		this.transformer = transformer;
+		this.beforeStream = beforeStream;
 		this.chunkTransformer = chunkTransformer;
 		this.afterChunkUpdate = afterChunkUpdate;
 	}
@@ -52,6 +55,9 @@ public class TextAgentResponseHandler implements StreamedAgentResponseHandler {
 	@Override
 	public void beforeStream(final AgentContext context) {
 		context.setStringValue(this.key, "");
+		if (this.beforeStream != null) {
+			this.beforeStream.accept(context);
+		}
 	}
 	
 	@Override
@@ -60,11 +66,13 @@ public class TextAgentResponseHandler implements StreamedAgentResponseHandler {
 		if (this.chunkTransformer != null) {
 			chunk = this.chunkTransformer.apply(chunk);
 		}
-		context.setStringValue(this.chunkKey, chunk);
-		context.setStringValue(this.key, Objects.requireNonNullElse(context.getStringValue(this.key), "").concat(chunk));
 		context.setStringValue(UPDATE_KEY, this.key);
 		if (afterChunkUpdate != null) {
+			context.setStringValue(this.chunkKey, chunk);
+			context.setStringValue(this.key, Objects.requireNonNullElse(context.getStringValue(this.key), "").concat(chunk));
 			afterChunkUpdate.accept(context);
+			context.getContext().remove(UPDATE_KEY);
+			context.getContext().remove(this.chunkKey);
 		}
 		return true;
 	}