From 7c8ff22753ddc05fbaff604b9ea3e4aa371fd5ad Mon Sep 17 00:00:00 2001
From: Guillaume Aquilina <guillaume@workflowai.com>
Date: Thu, 10 Apr 2025 10:31:51 +0200
Subject: [PATCH] feat: add tenant uid to agent

---
 workflowai/core/client/_models.py    | 1 +
 workflowai/core/client/_types.py     | 3 +++
 workflowai/core/client/agent.py      | 1 +
 workflowai/core/client/agent_test.py | 3 ++-
 4 files changed, 7 insertions(+), 1 deletion(-)

diff --git a/workflowai/core/client/_models.py b/workflowai/core/client/_models.py
index b73cc73..6da5a3e 100644
--- a/workflowai/core/client/_models.py
+++ b/workflowai/core/client/_models.py
@@ -167,6 +167,7 @@ class CreateAgentResponse(BaseModel):
     id: str
     schema_id: int
     uid: int = 0
+    tenant_uid: int = 0
 
 
 class ModelMetadata(BaseModel):
diff --git a/workflowai/core/client/_types.py b/workflowai/core/client/_types.py
index 7771c49..1635c0c 100644
--- a/workflowai/core/client/_types.py
+++ b/workflowai/core/client/_types.py
@@ -57,6 +57,9 @@ class AgentInterface(_BaseObject, Generic[AgentInputContra, AgentOutput], Protoc
     @property
     def agent_uid(self) -> int: ...
 
+    @property
+    def tenant_uid(self) -> int: ...
+
     async def run(
         self,
         agent_input: AgentInputContra,
diff --git a/workflowai/core/client/agent.py b/workflowai/core/client/agent.py
index 7a8c86d..07b321d 100644
--- a/workflowai/core/client/agent.py
+++ b/workflowai/core/client/agent.py
@@ -270,6 +270,7 @@ async def register(self):
         )
         self.schema_id = res.schema_id
         self.agent_uid = res.uid
+        self.tenant_uid = res.tenant_uid
         return res.schema_id
 
     @classmethod
diff --git a/workflowai/core/client/agent_test.py b/workflowai/core/client/agent_test.py
index 031cda8..b543803 100644
--- a/workflowai/core/client/agent_test.py
+++ b/workflowai/core/client/agent_test.py
@@ -232,6 +232,7 @@ async def test_auto_register(self, httpx_mock: HTTPXMock, agent_no_schema: Agent
                 "id": "123",
                 "schema_id": 2,
                 "uid": 123,
+                "tenant_uid": 1234,
             },
         )
         run_response = fixtures_json("task_run.json")
@@ -243,7 +244,7 @@ async def test_auto_register(self, httpx_mock: HTTPXMock, agent_no_schema: Agent
         out = await agent_no_schema.run(HelloTaskInput(name="Alice"))
         assert out.id == "8f635b73-f403-47ee-bff9-18320616c6cc"
         assert agent_no_schema.agent_uid == 123
-
+        assert agent_no_schema.tenant_uid == 1234
         run_response["id"] = "8f635b73-f403-47ee-bff9-18320616c6cc"
         # Try and run again
         httpx_mock.add_response(