Updates to proto for state apis (#5407)

2025-06-26 22:30:10 +00:00 · 2025-02-06 16:54:21 -05:00 · 2025-02-06 16:54:21 -05:00 · 25f26a338b
commit 25f26a338b
parent da6f918708
7 changed files with 202 additions and 188 deletions
--- a/dotnet/test/Microsoft.AutoGen.Core.Grpc.Tests/GrpcAgentServiceFixture.cs
+++ b/dotnet/test/Microsoft.AutoGen.Core.Grpc.Tests/GrpcAgentServiceFixture.cs
@ -25,8 +25,6 @@ public sealed class GrpcAgentServiceFixture() : AgentRpc.AgentRpcBase
            throw;
        }
    }
-    public override async Task<GetStateResponse> GetState(AgentId request, ServerCallContext context) => new GetStateResponse { AgentState = new AgentState { AgentId = request } };
-    public override async Task<SaveStateResponse> SaveState(AgentState request, ServerCallContext context) => new SaveStateResponse { };
    public override async Task<AddSubscriptionResponse> AddSubscription(AddSubscriptionRequest request, ServerCallContext context) => new AddSubscriptionResponse { };
    public override async Task<RemoveSubscriptionResponse> RemoveSubscription(RemoveSubscriptionRequest request, ServerCallContext context) => new RemoveSubscriptionResponse { };
    public override async Task<GetSubscriptionsResponse> GetSubscriptions(GetSubscriptionsRequest request, ServerCallContext context) => new GetSubscriptionsResponse { };
--- a/protos/agent_worker.proto
+++ b/protos/agent_worker.proto
@ -79,23 +79,6 @@ message GetSubscriptionsResponse {
    repeated Subscription subscriptions = 1;
 }

-message AgentState {
-  AgentId agent_id = 1;
-  string eTag = 2;
-  oneof  data {
-    bytes binary_data = 3;
-    string text_data = 4;
-    google.protobuf.Any proto_data = 5;
-  }
-}
-
-message GetStateResponse {
-	AgentState agent_state = 1;
-}
-
-message SaveStateResponse {
-}
-
 message Message {
    oneof message {
        RpcRequest request = 1;
@ -104,10 +87,46 @@ message Message {
    }
 }

+message SaveStateRequest {
+    AgentId agentId = 1;
+}
+
+message SaveStateResponse {
+    string state = 1;
+    optional string error = 2;
+}
+
+message LoadStateRequest {
+    AgentId agentId = 1;
+    string state = 2;
+}
+message LoadStateResponse {
+    optional string error = 1;
+}
+
+message ControlMessage {
+    // A response message should have the same id as the request message
+    string rpc_id = 1;
+    // This is either:
+    // agentid=AGENT_ID
+    // clientid=CLIENT_ID
+    string destination = 2;
+    // This is either:
+    // agentid=AGENT_ID
+    // clientid=CLIENT_ID
+    // Empty string means the message is a response
+    optional string respond_to = 3;
+    // One of:
+    //     SaveStateRequest saveStateRequest = 2;
+    //     SaveStateResponse saveStateResponse = 3;
+    //     LoadStateRequest loadStateRequest = 4;
+    //     LoadStateResponse loadStateResponse = 5;
+    google.protobuf.Any rpcMessage = 4;
+}
+
 service AgentRpc {
    rpc OpenChannel (stream Message) returns (stream Message);
-    rpc GetState(AgentId) returns (GetStateResponse);
-    rpc SaveState(AgentState) returns (SaveStateResponse);
+    rpc OpenControlChannel (stream ControlMessage) returns (stream ControlMessage);
    rpc RegisterAgent(RegisterAgentTypeRequest) returns (RegisterAgentTypeResponse);
    rpc AddSubscription(AddSubscriptionRequest) returns (AddSubscriptionResponse);
    rpc RemoveSubscription(RemoveSubscriptionRequest) returns (RemoveSubscriptionResponse);
--- a/python/packages/autogen-ext/src/autogen_ext/runtimes/grpc/_worker_runtime_host_servicer.py
+++ b/python/packages/autogen-ext/src/autogen_ext/runtimes/grpc/_worker_runtime_host_servicer.py
@ -135,6 +135,13 @@ class GrpcWorkerAgentRuntimeHostServicer(agent_worker_pb2_grpc.AgentRpcServicer)
            # Remove the client id from the agent type to client id mapping.
            await self._on_client_disconnect(client_id)

+    async def OpenControlChannel(  # type: ignore
+        self,
+        request_iterator: AsyncIterator[agent_worker_pb2.ControlMessage],
+        context: grpc.aio.ServicerContext[agent_worker_pb2.ControlMessage, agent_worker_pb2.ControlMessage],
+    ) -> AsyncIterator[agent_worker_pb2.ControlMessage]:
+        raise NotImplementedError("Method not implemented.")
+
    async def _on_client_disconnect(self, client_id: ClientConnectionId) -> None:
        async with self._agent_type_to_client_id_lock:
            agent_types = [agent_type for agent_type, id_ in self._agent_type_to_client_id.items() if id_ == client_id]
@ -288,17 +295,3 @@ class GrpcWorkerAgentRuntimeHostServicer(agent_worker_pb2_grpc.AgentRpcServicer)
    ) -> agent_worker_pb2.GetSubscriptionsResponse:
        _client_id = await get_client_id_or_abort(context)
        raise NotImplementedError("Method not implemented.")
-
-    async def GetState(  # type: ignore
-        self,
-        request: agent_worker_pb2.AgentId,
-        context: grpc.aio.ServicerContext[agent_worker_pb2.AgentId, agent_worker_pb2.GetStateResponse],
-    ) -> agent_worker_pb2.GetStateResponse:
-        raise NotImplementedError("Method not implemented!")
-
-    async def SaveState(  # type: ignore
-        self,
-        request: agent_worker_pb2.AgentState,
-        context: grpc.aio.ServicerContext[agent_worker_pb2.AgentId, agent_worker_pb2.SaveStateResponse],
-    ) -> agent_worker_pb2.SaveStateResponse:
-        raise NotImplementedError("Method not implemented!")
--- a/python/packages/autogen-ext/src/autogen_ext/runtimes/grpc/protos/agent_worker_pb2.py
+++ b/python/packages/autogen-ext/src/autogen_ext/runtimes/grpc/protos/agent_worker_pb2.py
@ -26,7 +26,7 @@ from . import cloudevent_pb2 as cloudevent__pb2
 from google.protobuf import any_pb2 as google_dot_protobuf_dot_any__pb2


-DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x12\x61gent_worker.proto\x12\x06\x61gents\x1a\x10\x63loudevent.proto\x1a\x19google/protobuf/any.proto\"$\n\x07\x41gentId\x12\x0c\n\x04type\x18\x01 \x01(\t\x12\x0b\n\x03key\x18\x02 \x01(\t\"E\n\x07Payload\x12\x11\n\tdata_type\x18\x01 \x01(\t\x12\x19\n\x11\x64\x61ta_content_type\x18\x02 \x01(\t\x12\x0c\n\x04\x64\x61ta\x18\x03 \x01(\x0c\"\x89\x02\n\nRpcRequest\x12\x12\n\nrequest_id\x18\x01 \x01(\t\x12$\n\x06source\x18\x02 \x01(\x0b\x32\x0f.agents.AgentIdH\x00\x88\x01\x01\x12\x1f\n\x06target\x18\x03 \x01(\x0b\x32\x0f.agents.AgentId\x12\x0e\n\x06method\x18\x04 \x01(\t\x12 \n\x07payload\x18\x05 \x01(\x0b\x32\x0f.agents.Payload\x12\x32\n\x08metadata\x18\x06 \x03(\x0b\x32 .agents.RpcRequest.MetadataEntry\x1a/\n\rMetadataEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\t:\x02\x38\x01\x42\t\n\x07_source\"\xb8\x01\n\x0bRpcResponse\x12\x12\n\nrequest_id\x18\x01 \x01(\t\x12 \n\x07payload\x18\x02 \x01(\x0b\x32\x0f.agents.Payload\x12\r\n\x05\x65rror\x18\x03 \x01(\t\x12\x33\n\x08metadata\x18\x04 \x03(\x0b\x32!.agents.RpcResponse.MetadataEntry\x1a/\n\rMetadataEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\t:\x02\x38\x01\"(\n\x18RegisterAgentTypeRequest\x12\x0c\n\x04type\x18\x01 \x01(\t\"\x1b\n\x19RegisterAgentTypeResponse\":\n\x10TypeSubscription\x12\x12\n\ntopic_type\x18\x01 \x01(\t\x12\x12\n\nagent_type\x18\x02 \x01(\t\"G\n\x16TypePrefixSubscription\x12\x19\n\x11topic_type_prefix\x18\x01 \x01(\t\x12\x12\n\nagent_type\x18\x02 \x01(\t\"\xa2\x01\n\x0cSubscription\x12\n\n\x02id\x18\x01 \x01(\t\x12\x34\n\x10typeSubscription\x18\x02 \x01(\x0b\x32\x18.agents.TypeSubscriptionH\x00\x12@\n\x16typePrefixSubscription\x18\x03 \x01(\x0b\x32\x1e.agents.TypePrefixSubscriptionH\x00\x42\x0e\n\x0csubscription\"D\n\x16\x41\x64\x64SubscriptionRequest\x12*\n\x0csubscription\x18\x01 \x01(\x0b\x32\x14.agents.Subscription\"\x19\n\x17\x41\x64\x64SubscriptionResponse\"\'\n\x19RemoveSubscriptionRequest\x12\n\n\x02id\x18\x01 \x01(\t\"\x1c\n\x1aRemoveSubscriptionResponse\"\x19\n\x17GetSubscriptionsRequest\"G\n\x18GetSubscriptionsResponse\x12+\n\rsubscriptions\x18\x01 \x03(\x0b\x32\x14.agents.Subscription\"\x9d\x01\n\nAgentState\x12!\n\x08\x61gent_id\x18\x01 \x01(\x0b\x32\x0f.agents.AgentId\x12\x0c\n\x04\x65Tag\x18\x02 \x01(\t\x12\x15\n\x0b\x62inary_data\x18\x03 \x01(\x0cH\x00\x12\x13\n\ttext_data\x18\x04 \x01(\tH\x00\x12*\n\nproto_data\x18\x05 \x01(\x0b\x32\x14.google.protobuf.AnyH\x00\x42\x06\n\x04\x64\x61ta\";\n\x10GetStateResponse\x12\'\n\x0b\x61gent_state\x18\x01 \x01(\x0b\x32\x12.agents.AgentState\"\x13\n\x11SaveStateResponse\"\x99\x01\n\x07Message\x12%\n\x07request\x18\x01 \x01(\x0b\x32\x12.agents.RpcRequestH\x00\x12\'\n\x08response\x18\x02 \x01(\x0b\x32\x13.agents.RpcResponseH\x00\x12\x33\n\ncloudEvent\x18\x03 \x01(\x0b\x32\x1d.io.cloudevents.v1.CloudEventH\x00\x42\t\n\x07message2\x90\x04\n\x08\x41gentRpc\x12\x33\n\x0bOpenChannel\x12\x0f.agents.Message\x1a\x0f.agents.Message(\x01\x30\x01\x12\x35\n\x08GetState\x12\x0f.agents.AgentId\x1a\x18.agents.GetStateResponse\x12:\n\tSaveState\x12\x12.agents.AgentState\x1a\x19.agents.SaveStateResponse\x12T\n\rRegisterAgent\x12 .agents.RegisterAgentTypeRequest\x1a!.agents.RegisterAgentTypeResponse\x12R\n\x0f\x41\x64\x64Subscription\x12\x1e.agents.AddSubscriptionRequest\x1a\x1f.agents.AddSubscriptionResponse\x12[\n\x12RemoveSubscription\x12!.agents.RemoveSubscriptionRequest\x1a\".agents.RemoveSubscriptionResponse\x12U\n\x10GetSubscriptions\x12\x1f.agents.GetSubscriptionsRequest\x1a .agents.GetSubscriptionsResponseB\x1d\xaa\x02\x1aMicrosoft.AutoGen.Protobufb\x06proto3')
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x12\x61gent_worker.proto\x12\x06\x61gents\x1a\x10\x63loudevent.proto\x1a\x19google/protobuf/any.proto\"$\n\x07\x41gentId\x12\x0c\n\x04type\x18\x01 \x01(\t\x12\x0b\n\x03key\x18\x02 \x01(\t\"E\n\x07Payload\x12\x11\n\tdata_type\x18\x01 \x01(\t\x12\x19\n\x11\x64\x61ta_content_type\x18\x02 \x01(\t\x12\x0c\n\x04\x64\x61ta\x18\x03 \x01(\x0c\"\x89\x02\n\nRpcRequest\x12\x12\n\nrequest_id\x18\x01 \x01(\t\x12$\n\x06source\x18\x02 \x01(\x0b\x32\x0f.agents.AgentIdH\x00\x88\x01\x01\x12\x1f\n\x06target\x18\x03 \x01(\x0b\x32\x0f.agents.AgentId\x12\x0e\n\x06method\x18\x04 \x01(\t\x12 \n\x07payload\x18\x05 \x01(\x0b\x32\x0f.agents.Payload\x12\x32\n\x08metadata\x18\x06 \x03(\x0b\x32 .agents.RpcRequest.MetadataEntry\x1a/\n\rMetadataEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\t:\x02\x38\x01\x42\t\n\x07_source\"\xb8\x01\n\x0bRpcResponse\x12\x12\n\nrequest_id\x18\x01 \x01(\t\x12 \n\x07payload\x18\x02 \x01(\x0b\x32\x0f.agents.Payload\x12\r\n\x05\x65rror\x18\x03 \x01(\t\x12\x33\n\x08metadata\x18\x04 \x03(\x0b\x32!.agents.RpcResponse.MetadataEntry\x1a/\n\rMetadataEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\t:\x02\x38\x01\"(\n\x18RegisterAgentTypeRequest\x12\x0c\n\x04type\x18\x01 \x01(\t\"\x1b\n\x19RegisterAgentTypeResponse\":\n\x10TypeSubscription\x12\x12\n\ntopic_type\x18\x01 \x01(\t\x12\x12\n\nagent_type\x18\x02 \x01(\t\"G\n\x16TypePrefixSubscription\x12\x19\n\x11topic_type_prefix\x18\x01 \x01(\t\x12\x12\n\nagent_type\x18\x02 \x01(\t\"\xa2\x01\n\x0cSubscription\x12\n\n\x02id\x18\x01 \x01(\t\x12\x34\n\x10typeSubscription\x18\x02 \x01(\x0b\x32\x18.agents.TypeSubscriptionH\x00\x12@\n\x16typePrefixSubscription\x18\x03 \x01(\x0b\x32\x1e.agents.TypePrefixSubscriptionH\x00\x42\x0e\n\x0csubscription\"D\n\x16\x41\x64\x64SubscriptionRequest\x12*\n\x0csubscription\x18\x01 \x01(\x0b\x32\x14.agents.Subscription\"\x19\n\x17\x41\x64\x64SubscriptionResponse\"\'\n\x19RemoveSubscriptionRequest\x12\n\n\x02id\x18\x01 \x01(\t\"\x1c\n\x1aRemoveSubscriptionResponse\"\x19\n\x17GetSubscriptionsRequest\"G\n\x18GetSubscriptionsResponse\x12+\n\rsubscriptions\x18\x01 \x03(\x0b\x32\x14.agents.Subscription\"\x99\x01\n\x07Message\x12%\n\x07request\x18\x01 \x01(\x0b\x32\x12.agents.RpcRequestH\x00\x12\'\n\x08response\x18\x02 \x01(\x0b\x32\x13.agents.RpcResponseH\x00\x12\x33\n\ncloudEvent\x18\x03 \x01(\x0b\x32\x1d.io.cloudevents.v1.CloudEventH\x00\x42\t\n\x07message\"4\n\x10SaveStateRequest\x12 \n\x07\x61gentId\x18\x01 \x01(\x0b\x32\x0f.agents.AgentId\"@\n\x11SaveStateResponse\x12\r\n\x05state\x18\x01 \x01(\t\x12\x12\n\x05\x65rror\x18\x02 \x01(\tH\x00\x88\x01\x01\x42\x08\n\x06_error\"C\n\x10LoadStateRequest\x12 \n\x07\x61gentId\x18\x01 \x01(\x0b\x32\x0f.agents.AgentId\x12\r\n\x05state\x18\x02 \x01(\t\"1\n\x11LoadStateResponse\x12\x12\n\x05\x65rror\x18\x01 \x01(\tH\x00\x88\x01\x01\x42\x08\n\x06_error\"\x87\x01\n\x0e\x43ontrolMessage\x12\x0e\n\x06rpc_id\x18\x01 \x01(\t\x12\x13\n\x0b\x64\x65stination\x18\x02 \x01(\t\x12\x17\n\nrespond_to\x18\x03 \x01(\tH\x00\x88\x01\x01\x12(\n\nrpcMessage\x18\x04 \x01(\x0b\x32\x14.google.protobuf.AnyB\r\n\x0b_respond_to2\xe7\x03\n\x08\x41gentRpc\x12\x33\n\x0bOpenChannel\x12\x0f.agents.Message\x1a\x0f.agents.Message(\x01\x30\x01\x12H\n\x12OpenControlChannel\x12\x16.agents.ControlMessage\x1a\x16.agents.ControlMessage(\x01\x30\x01\x12T\n\rRegisterAgent\x12 .agents.RegisterAgentTypeRequest\x1a!.agents.RegisterAgentTypeResponse\x12R\n\x0f\x41\x64\x64Subscription\x12\x1e.agents.AddSubscriptionRequest\x1a\x1f.agents.AddSubscriptionResponse\x12[\n\x12RemoveSubscription\x12!.agents.RemoveSubscriptionRequest\x1a\".agents.RemoveSubscriptionResponse\x12U\n\x10GetSubscriptions\x12\x1f.agents.GetSubscriptionsRequest\x1a .agents.GetSubscriptionsResponseB\x1d\xaa\x02\x1aMicrosoft.AutoGen.Protobufb\x06proto3')

 _globals = globals()
 _builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
@ -72,14 +72,18 @@ if not _descriptor._USE_C_DESCRIPTORS:
  _globals['_GETSUBSCRIPTIONSREQUEST']._serialized_end=1201
  _globals['_GETSUBSCRIPTIONSRESPONSE']._serialized_start=1203
  _globals['_GETSUBSCRIPTIONSRESPONSE']._serialized_end=1274
-  _globals['_AGENTSTATE']._serialized_start=1277
-  _globals['_AGENTSTATE']._serialized_end=1434
-  _globals['_GETSTATERESPONSE']._serialized_start=1436
-  _globals['_GETSTATERESPONSE']._serialized_end=1495
-  _globals['_SAVESTATERESPONSE']._serialized_start=1497
-  _globals['_SAVESTATERESPONSE']._serialized_end=1516
-  _globals['_MESSAGE']._serialized_start=1519
-  _globals['_MESSAGE']._serialized_end=1672
-  _globals['_AGENTRPC']._serialized_start=1675
-  _globals['_AGENTRPC']._serialized_end=2203
+  _globals['_MESSAGE']._serialized_start=1277
+  _globals['_MESSAGE']._serialized_end=1430
+  _globals['_SAVESTATEREQUEST']._serialized_start=1432
+  _globals['_SAVESTATEREQUEST']._serialized_end=1484
+  _globals['_SAVESTATERESPONSE']._serialized_start=1486
+  _globals['_SAVESTATERESPONSE']._serialized_end=1550
+  _globals['_LOADSTATEREQUEST']._serialized_start=1552
+  _globals['_LOADSTATEREQUEST']._serialized_end=1619
+  _globals['_LOADSTATERESPONSE']._serialized_start=1621
+  _globals['_LOADSTATERESPONSE']._serialized_end=1670
+  _globals['_CONTROLMESSAGE']._serialized_start=1673
+  _globals['_CONTROLMESSAGE']._serialized_end=1808
+  _globals['_AGENTRPC']._serialized_start=1811
+  _globals['_AGENTRPC']._serialized_end=2298
 # @@protoc_insertion_point(module_scope)
--- a/python/packages/autogen-ext/src/autogen_ext/runtimes/grpc/protos/agent_worker_pb2.pyi
+++ b/python/packages/autogen-ext/src/autogen_ext/runtimes/grpc/protos/agent_worker_pb2.pyi
@ -312,64 +312,6 @@ class GetSubscriptionsResponse(google.protobuf.message.Message):

 global___GetSubscriptionsResponse = GetSubscriptionsResponse

-@typing.final
-class AgentState(google.protobuf.message.Message):
-    DESCRIPTOR: google.protobuf.descriptor.Descriptor
-
-    AGENT_ID_FIELD_NUMBER: builtins.int
-    ETAG_FIELD_NUMBER: builtins.int
-    BINARY_DATA_FIELD_NUMBER: builtins.int
-    TEXT_DATA_FIELD_NUMBER: builtins.int
-    PROTO_DATA_FIELD_NUMBER: builtins.int
-    eTag: builtins.str
-    binary_data: builtins.bytes
-    text_data: builtins.str
-    @property
-    def agent_id(self) -> global___AgentId: ...
-    @property
-    def proto_data(self) -> google.protobuf.any_pb2.Any: ...
-    def __init__(
-        self,
-        *,
-        agent_id: global___AgentId | None = ...,
-        eTag: builtins.str = ...,
-        binary_data: builtins.bytes = ...,
-        text_data: builtins.str = ...,
-        proto_data: google.protobuf.any_pb2.Any | None = ...,
-    ) -> None: ...
-    def HasField(self, field_name: typing.Literal["agent_id", b"agent_id", "binary_data", b"binary_data", "data", b"data", "proto_data", b"proto_data", "text_data", b"text_data"]) -> builtins.bool: ...
-    def ClearField(self, field_name: typing.Literal["agent_id", b"agent_id", "binary_data", b"binary_data", "data", b"data", "eTag", b"eTag", "proto_data", b"proto_data", "text_data", b"text_data"]) -> None: ...
-    def WhichOneof(self, oneof_group: typing.Literal["data", b"data"]) -> typing.Literal["binary_data", "text_data", "proto_data"] | None: ...
-
-global___AgentState = AgentState
-
-@typing.final
-class GetStateResponse(google.protobuf.message.Message):
-    DESCRIPTOR: google.protobuf.descriptor.Descriptor
-
-    AGENT_STATE_FIELD_NUMBER: builtins.int
-    @property
-    def agent_state(self) -> global___AgentState: ...
-    def __init__(
-        self,
-        *,
-        agent_state: global___AgentState | None = ...,
-    ) -> None: ...
-    def HasField(self, field_name: typing.Literal["agent_state", b"agent_state"]) -> builtins.bool: ...
-    def ClearField(self, field_name: typing.Literal["agent_state", b"agent_state"]) -> None: ...
-
-global___GetStateResponse = GetStateResponse
-
-@typing.final
-class SaveStateResponse(google.protobuf.message.Message):
-    DESCRIPTOR: google.protobuf.descriptor.Descriptor
-
-    def __init__(
-        self,
-    ) -> None: ...
-
-global___SaveStateResponse = SaveStateResponse
-
@typing.final
 class Message(google.protobuf.message.Message):
    DESCRIPTOR: google.protobuf.descriptor.Descriptor
@ -395,3 +337,121 @@ class Message(google.protobuf.message.Message):
    def WhichOneof(self, oneof_group: typing.Literal["message", b"message"]) -> typing.Literal["request", "response", "cloudEvent"] | None: ...

 global___Message = Message
+
+@typing.final
+class SaveStateRequest(google.protobuf.message.Message):
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+    AGENTID_FIELD_NUMBER: builtins.int
+    @property
+    def agentId(self) -> global___AgentId: ...
+    def __init__(
+        self,
+        *,
+        agentId: global___AgentId | None = ...,
+    ) -> None: ...
+    def HasField(self, field_name: typing.Literal["agentId", b"agentId"]) -> builtins.bool: ...
+    def ClearField(self, field_name: typing.Literal["agentId", b"agentId"]) -> None: ...
+
+global___SaveStateRequest = SaveStateRequest
+
+@typing.final
+class SaveStateResponse(google.protobuf.message.Message):
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+    STATE_FIELD_NUMBER: builtins.int
+    ERROR_FIELD_NUMBER: builtins.int
+    state: builtins.str
+    error: builtins.str
+    def __init__(
+        self,
+        *,
+        state: builtins.str = ...,
+        error: builtins.str | None = ...,
+    ) -> None: ...
+    def HasField(self, field_name: typing.Literal["_error", b"_error", "error", b"error"]) -> builtins.bool: ...
+    def ClearField(self, field_name: typing.Literal["_error", b"_error", "error", b"error", "state", b"state"]) -> None: ...
+    def WhichOneof(self, oneof_group: typing.Literal["_error", b"_error"]) -> typing.Literal["error"] | None: ...
+
+global___SaveStateResponse = SaveStateResponse
+
+@typing.final
+class LoadStateRequest(google.protobuf.message.Message):
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+    AGENTID_FIELD_NUMBER: builtins.int
+    STATE_FIELD_NUMBER: builtins.int
+    state: builtins.str
+    @property
+    def agentId(self) -> global___AgentId: ...
+    def __init__(
+        self,
+        *,
+        agentId: global___AgentId | None = ...,
+        state: builtins.str = ...,
+    ) -> None: ...
+    def HasField(self, field_name: typing.Literal["agentId", b"agentId"]) -> builtins.bool: ...
+    def ClearField(self, field_name: typing.Literal["agentId", b"agentId", "state", b"state"]) -> None: ...
+
+global___LoadStateRequest = LoadStateRequest
+
+@typing.final
+class LoadStateResponse(google.protobuf.message.Message):
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+    ERROR_FIELD_NUMBER: builtins.int
+    error: builtins.str
+    def __init__(
+        self,
+        *,
+        error: builtins.str | None = ...,
+    ) -> None: ...
+    def HasField(self, field_name: typing.Literal["_error", b"_error", "error", b"error"]) -> builtins.bool: ...
+    def ClearField(self, field_name: typing.Literal["_error", b"_error", "error", b"error"]) -> None: ...
+    def WhichOneof(self, oneof_group: typing.Literal["_error", b"_error"]) -> typing.Literal["error"] | None: ...
+
+global___LoadStateResponse = LoadStateResponse
+
+@typing.final
+class ControlMessage(google.protobuf.message.Message):
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+    RPC_ID_FIELD_NUMBER: builtins.int
+    DESTINATION_FIELD_NUMBER: builtins.int
+    RESPOND_TO_FIELD_NUMBER: builtins.int
+    RPCMESSAGE_FIELD_NUMBER: builtins.int
+    rpc_id: builtins.str
+    """A response message should have the same id as the request message"""
+    destination: builtins.str
+    """This is either:
+    agentid=AGENT_ID
+    clientid=CLIENT_ID
+    """
+    respond_to: builtins.str
+    """This is either:
+    agentid=AGENT_ID
+    clientid=CLIENT_ID
+    Empty string means the message is a response
+    """
+    @property
+    def rpcMessage(self) -> google.protobuf.any_pb2.Any:
+        """One of:
+            SaveStateRequest saveStateRequest = 2;
+            SaveStateResponse saveStateResponse = 3;
+            LoadStateRequest loadStateRequest = 4;
+            LoadStateResponse loadStateResponse = 5;
+        """
+
+    def __init__(
+        self,
+        *,
+        rpc_id: builtins.str = ...,
+        destination: builtins.str = ...,
+        respond_to: builtins.str | None = ...,
+        rpcMessage: google.protobuf.any_pb2.Any | None = ...,
+    ) -> None: ...
+    def HasField(self, field_name: typing.Literal["_respond_to", b"_respond_to", "respond_to", b"respond_to", "rpcMessage", b"rpcMessage"]) -> builtins.bool: ...
+    def ClearField(self, field_name: typing.Literal["_respond_to", b"_respond_to", "destination", b"destination", "respond_to", b"respond_to", "rpcMessage", b"rpcMessage", "rpc_id", b"rpc_id"]) -> None: ...
+    def WhichOneof(self, oneof_group: typing.Literal["_respond_to", b"_respond_to"]) -> typing.Literal["respond_to"] | None: ...
+
+global___ControlMessage = ControlMessage
--- a/python/packages/autogen-ext/src/autogen_ext/runtimes/grpc/protos/agent_worker_pb2_grpc.py
+++ b/python/packages/autogen-ext/src/autogen_ext/runtimes/grpc/protos/agent_worker_pb2_grpc.py
@ -39,15 +39,10 @@ class AgentRpcStub(object):
                request_serializer=agent__worker__pb2.Message.SerializeToString,
                response_deserializer=agent__worker__pb2.Message.FromString,
                _registered_method=True)
-        self.GetState = channel.unary_unary(
-                '/agents.AgentRpc/GetState',
-                request_serializer=agent__worker__pb2.AgentId.SerializeToString,
-                response_deserializer=agent__worker__pb2.GetStateResponse.FromString,
-                _registered_method=True)
-        self.SaveState = channel.unary_unary(
-                '/agents.AgentRpc/SaveState',
-                request_serializer=agent__worker__pb2.AgentState.SerializeToString,
-                response_deserializer=agent__worker__pb2.SaveStateResponse.FromString,
+        self.OpenControlChannel = channel.stream_stream(
+                '/agents.AgentRpc/OpenControlChannel',
+                request_serializer=agent__worker__pb2.ControlMessage.SerializeToString,
+                response_deserializer=agent__worker__pb2.ControlMessage.FromString,
                _registered_method=True)
        self.RegisterAgent = channel.unary_unary(
                '/agents.AgentRpc/RegisterAgent',
@ -80,13 +75,7 @@ class AgentRpcServicer(object):
        context.set_details('Method not implemented!')
        raise NotImplementedError('Method not implemented!')

-    def GetState(self, request, context):
-        """Missing associated documentation comment in .proto file."""
-        context.set_code(grpc.StatusCode.UNIMPLEMENTED)
-        context.set_details('Method not implemented!')
-        raise NotImplementedError('Method not implemented!')
-
-    def SaveState(self, request, context):
+    def OpenControlChannel(self, request_iterator, context):
        """Missing associated documentation comment in .proto file."""
        context.set_code(grpc.StatusCode.UNIMPLEMENTED)
        context.set_details('Method not implemented!')
@ -124,15 +113,10 @@ def add_AgentRpcServicer_to_server(servicer, server):
                    request_deserializer=agent__worker__pb2.Message.FromString,
                    response_serializer=agent__worker__pb2.Message.SerializeToString,
            ),
-            'GetState': grpc.unary_unary_rpc_method_handler(
-                    servicer.GetState,
-                    request_deserializer=agent__worker__pb2.AgentId.FromString,
-                    response_serializer=agent__worker__pb2.GetStateResponse.SerializeToString,
-            ),
-            'SaveState': grpc.unary_unary_rpc_method_handler(
-                    servicer.SaveState,
-                    request_deserializer=agent__worker__pb2.AgentState.FromString,
-                    response_serializer=agent__worker__pb2.SaveStateResponse.SerializeToString,
+            'OpenControlChannel': grpc.stream_stream_rpc_method_handler(
+                    servicer.OpenControlChannel,
+                    request_deserializer=agent__worker__pb2.ControlMessage.FromString,
+                    response_serializer=agent__worker__pb2.ControlMessage.SerializeToString,
            ),
            'RegisterAgent': grpc.unary_unary_rpc_method_handler(
                    servicer.RegisterAgent,
@ -193,7 +177,7 @@ class AgentRpc(object):
            _registered_method=True)

    @staticmethod
-    def GetState(request,
+    def OpenControlChannel(request_iterator,
            target,
            options=(),
            channel_credentials=None,
@ -203,39 +187,12 @@ class AgentRpc(object):
            wait_for_ready=None,
            timeout=None,
            metadata=None):
-        return grpc.experimental.unary_unary(
-            request,
+        return grpc.experimental.stream_stream(
+            request_iterator,
            target,
-            '/agents.AgentRpc/GetState',
-            agent__worker__pb2.AgentId.SerializeToString,
-            agent__worker__pb2.GetStateResponse.FromString,
-            options,
-            channel_credentials,
-            insecure,
-            call_credentials,
-            compression,
-            wait_for_ready,
-            timeout,
-            metadata,
-            _registered_method=True)
-
-    @staticmethod
-    def SaveState(request,
-            target,
-            options=(),
-            channel_credentials=None,
-            call_credentials=None,
-            insecure=False,
-            compression=None,
-            wait_for_ready=None,
-            timeout=None,
-            metadata=None):
-        return grpc.experimental.unary_unary(
-            request,
-            target,
-            '/agents.AgentRpc/SaveState',
-            agent__worker__pb2.AgentState.SerializeToString,
-            agent__worker__pb2.SaveStateResponse.FromString,
+            '/agents.AgentRpc/OpenControlChannel',
+            agent__worker__pb2.ControlMessage.SerializeToString,
+            agent__worker__pb2.ControlMessage.FromString,
            options,
            channel_credentials,
            insecure,
--- a/python/packages/autogen-ext/src/autogen_ext/runtimes/grpc/protos/agent_worker_pb2_grpc.pyi
+++ b/python/packages/autogen-ext/src/autogen_ext/runtimes/grpc/protos/agent_worker_pb2_grpc.pyi
@ -24,14 +24,9 @@ class AgentRpcStub:
        agent_worker_pb2.Message,
    ]

-    GetState: grpc.UnaryUnaryMultiCallable[
-        agent_worker_pb2.AgentId,
-        agent_worker_pb2.GetStateResponse,
-    ]
-
-    SaveState: grpc.UnaryUnaryMultiCallable[
-        agent_worker_pb2.AgentState,
-        agent_worker_pb2.SaveStateResponse,
+    OpenControlChannel: grpc.StreamStreamMultiCallable[
+        agent_worker_pb2.ControlMessage,
+        agent_worker_pb2.ControlMessage,
    ]

    RegisterAgent: grpc.UnaryUnaryMultiCallable[
@ -60,14 +55,9 @@ class AgentRpcAsyncStub:
        agent_worker_pb2.Message,
    ]

-    GetState: grpc.aio.UnaryUnaryMultiCallable[
-        agent_worker_pb2.AgentId,
-        agent_worker_pb2.GetStateResponse,
-    ]
-
-    SaveState: grpc.aio.UnaryUnaryMultiCallable[
-        agent_worker_pb2.AgentState,
-        agent_worker_pb2.SaveStateResponse,
+    OpenControlChannel: grpc.aio.StreamStreamMultiCallable[
+        agent_worker_pb2.ControlMessage,
+        agent_worker_pb2.ControlMessage,
    ]

    RegisterAgent: grpc.aio.UnaryUnaryMultiCallable[
@ -99,18 +89,11 @@ class AgentRpcServicer(metaclass=abc.ABCMeta):
    ) -> typing.Union[collections.abc.Iterator[agent_worker_pb2.Message], collections.abc.AsyncIterator[agent_worker_pb2.Message]]: ...

    @abc.abstractmethod
-    def GetState(
+    def OpenControlChannel(
        self,
-        request: agent_worker_pb2.AgentId,
+        request_iterator: _MaybeAsyncIterator[agent_worker_pb2.ControlMessage],
        context: _ServicerContext,
-    ) -> typing.Union[agent_worker_pb2.GetStateResponse, collections.abc.Awaitable[agent_worker_pb2.GetStateResponse]]: ...
-
-    @abc.abstractmethod
-    def SaveState(
-        self,
-        request: agent_worker_pb2.AgentState,
-        context: _ServicerContext,
-    ) -> typing.Union[agent_worker_pb2.SaveStateResponse, collections.abc.Awaitable[agent_worker_pb2.SaveStateResponse]]: ...
+    ) -> typing.Union[collections.abc.Iterator[agent_worker_pb2.ControlMessage], collections.abc.AsyncIterator[agent_worker_pb2.ControlMessage]]: ...

    @abc.abstractmethod
    def RegisterAgent(