From b028b826f0fa07538ec9ad15d1bc4318ec8f229c Mon Sep 17 00:00:00 2001
From: Christophe Bornet <cbornet@hotmail.com>
Date: Wed, 13 Dec 2023 16:33:00 +0100
Subject: [PATCH] Fix for handling big messages from Python agent

---
 .../main/java/ai/langstream/agents/grpc/GrpcAgentProcessor.java | 2 +-
 .../src/main/java/ai/langstream/agents/grpc/GrpcAgentSink.java  | 2 +-
 .../main/java/ai/langstream/agents/grpc/GrpcAgentSource.java    | 2 +-
 3 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/langstream-agents/langstream-agent-grpc/src/main/java/ai/langstream/agents/grpc/GrpcAgentProcessor.java b/langstream-agents/langstream-agent-grpc/src/main/java/ai/langstream/agents/grpc/GrpcAgentProcessor.java
index 528d9c79c..9ee9433dd 100644
--- a/langstream-agents/langstream-agent-grpc/src/main/java/ai/langstream/agents/grpc/GrpcAgentProcessor.java
+++ b/langstream-agents/langstream-agent-grpc/src/main/java/ai/langstream/agents/grpc/GrpcAgentProcessor.java
@@ -59,7 +59,7 @@ public synchronized void onNewSchemaToSend(Schema schema) {
     @Override
     public void start() throws Exception {
         super.start();
-        request = AgentServiceGrpc.newStub(channel).withWaitForReady().process(responseObserver);
+        request = asyncStub.process(responseObserver);
         restarting.set(false);
         startFailedButDevelopmentMode = false;
     }
diff --git a/langstream-agents/langstream-agent-grpc/src/main/java/ai/langstream/agents/grpc/GrpcAgentSink.java b/langstream-agents/langstream-agent-grpc/src/main/java/ai/langstream/agents/grpc/GrpcAgentSink.java
index c96fc9895..62f50e325 100644
--- a/langstream-agents/langstream-agent-grpc/src/main/java/ai/langstream/agents/grpc/GrpcAgentSink.java
+++ b/langstream-agents/langstream-agent-grpc/src/main/java/ai/langstream/agents/grpc/GrpcAgentSink.java
@@ -53,7 +53,7 @@ public void onNewSchemaToSend(Schema schema) {
     @Override
     public void start() throws Exception {
         super.start();
-        request = AgentServiceGrpc.newStub(channel).withWaitForReady().write(responseObserver);
+        request = asyncStub.write(responseObserver);
         restarting.set(false);
         startFailedButDevelopmentMode = false;
     }
diff --git a/langstream-agents/langstream-agent-grpc/src/main/java/ai/langstream/agents/grpc/GrpcAgentSource.java b/langstream-agents/langstream-agent-grpc/src/main/java/ai/langstream/agents/grpc/GrpcAgentSource.java
index 30a75ae41..65f672ab2 100644
--- a/langstream-agents/langstream-agent-grpc/src/main/java/ai/langstream/agents/grpc/GrpcAgentSource.java
+++ b/langstream-agents/langstream-agent-grpc/src/main/java/ai/langstream/agents/grpc/GrpcAgentSource.java
@@ -53,7 +53,7 @@ public void onNewSchemaToSend(Schema schema) {
     @Override
     public void start() throws Exception {
         super.start();
-        request = AgentServiceGrpc.newStub(channel).withWaitForReady().read(responseObserver);
+        request = asyncStub.read(responseObserver);
         restarting.set(false);
         startFailedButDevelopmentMode = false;
     }