Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
23 changes: 16 additions & 7 deletions sdk-python/agentcube/clients/agent_runtime_data_plane.py
Original file line number Diff line number Diff line change
Expand Up @@ -56,14 +56,23 @@ def bootstrap_session_id(self) -> str:
self.base_url,
timeout=(self.connect_timeout, self.timeout),
)
resp.raise_for_status()

session_id = resp.headers.get(self.SESSION_HEADER)
if not session_id:
raise ValueError(
f"Missing required response header: {self.SESSION_HEADER}"
)
return session_id
if session_id:
if resp.status_code >= 400:
self.logger.debug(
f"Bootstrap request returned status {resp.status_code}, "
f"but session ID was found: {session_id}"
)
Comment on lines +61 to +65
return session_id
Comment thread
Abhinav-kodes marked this conversation as resolved.
resp.raise_for_status()
content_type = resp.headers.get("Content-Type")
content_length = resp.headers.get("Content-Length")
raise ValueError(
f"Missing required response header: {self.SESSION_HEADER} "
f"(status: {resp.status_code}, "
f"content-type: {content_type}, "
f"content-length: {content_length})"
)
Comment thread
Abhinav-kodes marked this conversation as resolved.

def invoke(
self,
Expand Down
4 changes: 2 additions & 2 deletions sdk-python/examples/agent_runtime_usage.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@
# first time: it will create a new pod
agent_client_v1 = AgentRuntimeClient(
agent_name="my-agent",
router_url="http://localhost:18081",
router_url="http://localhost:8081",
namespace="default",
verbose=True,
)
Expand All @@ -31,7 +31,7 @@
# second time: it will try to reuse the pod created before
agent_client_v2 = AgentRuntimeClient(
agent_name="my-agent",
router_url="http://localhost:18081",
router_url="http://localhost:8081",
namespace="default",
session_id=agent_client_v1.session_id,
verbose=True,
Expand Down
21 changes: 21 additions & 0 deletions sdk-python/tests/test_agent_runtime.py
Original file line number Diff line number Diff line change
Expand Up @@ -95,6 +95,7 @@ class TestAgentRuntimeDataPlaneClient(unittest.TestCase):
def test_bootstrap_session_id_extracts_header(self, mock_create_session):
sess = Mock()
resp = Mock()
resp.status_code = 200
resp.raise_for_status.return_value = None
resp.headers = {"x-agentcube-session-id": "abc"}
sess.get.return_value = resp
Comment on lines 95 to 101
Expand All @@ -109,6 +110,26 @@ def test_bootstrap_session_id_extracts_header(self, mock_create_session):
)
self.assertEqual(client.bootstrap_session_id(), "abc")

@patch("agentcube.clients.agent_runtime_data_plane.create_session")
def test_bootstrap_session_id_returns_header_on_non_2xx(self, mock_create_session):
sess = Mock()
resp = Mock()
resp.status_code = 404
resp.raise_for_status.side_effect = requests.exceptions.HTTPError(response=resp)
resp.headers = {"x-agentcube-session-id": "session-from-404"}
sess.get.return_value = resp
mock_create_session.return_value = sess

from agentcube.clients.agent_runtime_data_plane import AgentRuntimeDataPlaneClient

client = AgentRuntimeDataPlaneClient(
router_url="http://router",
namespace="default",
agent_name="agent-a",
)
self.assertEqual(client.bootstrap_session_id(), "session-from-404")
resp.raise_for_status.assert_not_called()

@patch("agentcube.clients.agent_runtime_data_plane.create_session")
def test_invoke_sends_session_header(self, mock_create_session):
sess = Mock()
Expand Down
Loading