From 2f26632ed29521d57f063bc908d4d9ba0f863cd2 Mon Sep 17 00:00:00 2001
From: Roger Johansson <roger@asynkron.se>
Date: Thu, 27 May 2021 17:32:30 +0200
Subject: [PATCH] update template

---
 src/Proto.Cluster.CodeGen/Template.cs | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/src/Proto.Cluster.CodeGen/Template.cs b/src/Proto.Cluster.CodeGen/Template.cs
index d323e8b61f..311a0d5e29 100644
--- a/src/Proto.Cluster.CodeGen/Template.cs
+++ b/src/Proto.Cluster.CodeGen/Template.cs
@@ -15,13 +15,14 @@ public static class Template
 using Proto;
 using Proto.Cluster;
 
-
 namespace {{CsNamespace}}
 {
     public static class GrainExtensions
     {
         {{#each Services}}
         public static {{Name}}Client Get{{Name}}(this Cluster cluster, string identity) => new(cluster, identity);
+
+        public static {{Name}}Client Get{{Name}}(this IContext context, string identity) => new(context.System.Cluster(), identity);
         {{/each}}
     }
 
@@ -29,6 +30,8 @@ public static class GrainExtensions
     public abstract class {{Name}}Base
     {
         protected IContext Context {get;}
+        protected ActorSystem System => Context.System;
+        protected Cluster Cluster => Context.System.Cluster();
     
         protected {{Name}}Base(IContext context)
         {
@@ -112,7 +115,7 @@ class {{Name}}Actor : IActor
     {
         private {{Name}}Base _inner;
         private IContext _context;
-        private Func<IContext, string, string, {{Name}}Base> _innerFactory;
+        private Func<IContext, string, string, {{Name}}Base> _innerFactory;        
     
         public {{Name}}Actor(Func<IContext, string, string, {{Name}}Base> innerFactory)
         {