Skip to content
Merged
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
11 changes: 6 additions & 5 deletions sdk/python/agentfield/memory_events.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
from datetime import datetime
from functools import wraps
from typing import Any, Callable, Dict, List, Optional, Union
from urllib.parse import urlencode

import websockets

Expand Down Expand Up @@ -152,16 +153,16 @@ async def connect(
ws_url = f"{self.base_url}/api/v1/memory/events/ws"

# Add query parameters for server-side filtering
query_params = []
query_params = {}
if patterns:
query_params.append(f"patterns={','.join(patterns)}")
query_params["patterns"] = ",".join(patterns)
if scope:
query_params.append(f"scope={scope}")
query_params["scope"] = scope
if scope_id:
query_params.append(f"scope_id={scope_id}")
query_params["scope_id"] = scope_id

if query_params:
ws_url += "?" + "&".join(query_params)
ws_url += "?" + urlencode(query_params)

self.websocket = await asyncio.wait_for(
websockets.connect(ws_url, **{_HEADERS_KWARG: headers}),
Expand Down
42 changes: 32 additions & 10 deletions sdk/python/tests/test_memory_events.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
import asyncio
import json
from types import SimpleNamespace
from urllib.parse import parse_qs, urlsplit

import websockets

Expand Down Expand Up @@ -161,7 +162,18 @@ async def get(self, url, params=None, headers=None, timeout=None):


@pytest.mark.asyncio
async def test_memory_event_client_connect_builds_ws_url(monkeypatch):
@pytest.mark.parametrize(
("patterns", "scope", "scope_id"),
[
(["cart.*", "order.*"], "session", "abc"),
(["user+plan.*"], "actor", "alice+trial@example.com"),
(["cart&total.*", "issue#*"], "session", "team&env=prod#1"),
(["user preferences.*", "用户.*"], "workflow", "会话 + 1"),
],
)
async def test_memory_event_client_connect_builds_ws_url(
monkeypatch, patterns, scope, scope_id
):
ctx = SimpleNamespace(to_headers=lambda: {"Authorization": "token"})
client = MemoryEventClient("http://agentfield", ctx)

Expand All @@ -174,7 +186,9 @@ def __init__(self):

async def fake_connect(url, **kwargs):
record["url"] = url
record["headers"] = kwargs.get("additional_headers") or kwargs.get("extra_headers")
record["headers"] = kwargs.get("additional_headers") or kwargs.get(
"extra_headers"
)
return DummyWebSocket()

async def fake_listen(self):
Expand All @@ -183,15 +197,19 @@ async def fake_listen(self):
monkeypatch.setattr("agentfield.memory_events.websockets.connect", fake_connect)
monkeypatch.setattr(MemoryEventClient, "_listen", fake_listen, raising=False)

await client.connect(
patterns=["cart.*", "order.*"], scope="session", scope_id="abc"
)
await client.connect(patterns=patterns, scope=scope, scope_id=scope_id)
await asyncio.sleep(0)

assert record["url"].startswith("ws://agentfield")
assert "patterns=cart.*,order.*" in record["url"]
assert "scope=session" in record["url"]
assert "scope_id=abc" in record["url"]
parsed_url = urlsplit(record["url"])
assert parsed_url.scheme == "ws"
assert parsed_url.netloc == "agentfield"
assert parsed_url.path == "/api/v1/memory/events/ws"
assert parsed_url.fragment == ""
assert parse_qs(parsed_url.query) == {
"patterns": [",".join(patterns)],
"scope": [scope],
"scope_id": [scope_id],
}
assert record["headers"] == {"Authorization": "token"}
assert listener_called.get("run") is True

Expand Down Expand Up @@ -225,6 +243,7 @@ def __init__(self):

async def fake_connect(url, **kwargs):
called_with.update(kwargs)
called_with["url"] = url
return DummyWebSocket()

monkeypatch.setattr("agentfield.memory_events.websockets.connect", fake_connect)
Expand All @@ -235,6 +254,7 @@ async def fake_connect(url, **kwargs):

assert _HEADERS_KWARG in called_with
assert called_with[_HEADERS_KWARG] == {"Authorization": "token"}
assert called_with["url"] == "ws://agentfield/api/v1/memory/events/ws"


@pytest.mark.asyncio
Expand All @@ -260,7 +280,9 @@ async def fake_reconnect(self):
# Give the background task a chance to start
await asyncio.sleep(0.05)

assert reconnect_started.is_set(), "reconnect should have been started in background"
assert reconnect_started.is_set(), (
"reconnect should have been started in background"
)
assert not client.is_listening


Expand Down
Loading