diff --git a/codegen/core/src/main/java/software/amazon/smithy/python/codegen/ClientGenerator.java b/codegen/core/src/main/java/software/amazon/smithy/python/codegen/ClientGenerator.java
index a1c8abbe0..8e8b4fb49 100644
--- a/codegen/core/src/main/java/software/amazon/smithy/python/codegen/ClientGenerator.java
+++ b/codegen/core/src/main/java/software/amazon/smithy/python/codegen/ClientGenerator.java
@@ -905,7 +905,7 @@ private void writeSharedOperationInit(PythonWriter writer, OperationShape operat
     }
 
     private void generateEventStreamOperation(PythonWriter writer, OperationShape operation) {
-        writer.pushState();
+        writer.pushState(new OperationSection(service, operation));
         writer.addDependency(SmithyPythonDependency.SMITHY_CORE);
         var operationSymbol = symbolProvider.toSymbol(operation);
         writer.putContext("operation", operationSymbol);