fix(cuda.core): avoid truncating graph queries

cuda_core/cuda/core/graph/_adjacency_set_proxy.pyx

@@ -144,48 +144,48 @@ cdef class _AdjacencySetCore:
         cdef cydriver.CUgraphNode c_node = as_cu(self._h_node)
         if c_node == NULL:
             return []
-        cdef cydriver.CUgraphNode buf[16]
-        cdef size_t count = 16
+        cdef cydriver.CUgraphNode stack_buf[16]
+        cdef cydriver.CUgraphNode* nodes
+        cdef size_t count = 0
         cdef size_t i
         with nogil:
-            HANDLE_RETURN(self._query_fn(c_node, buf, &count))
-        if count <= 16:
-            return [GraphNode._create(self._h_graph, buf[i])
-                    for i in range(count)]
+            HANDLE_RETURN(self._query_fn(c_node, NULL, &count))
+        if count == 0:
+            return []
         cdef vector[cydriver.CUgraphNode] nodes_vec
-        nodes_vec.resize(count)
+        if count <= 16:
+            nodes = stack_buf
+        else:
+            nodes_vec.resize(count)
+            nodes = nodes_vec.data()
         with nogil:
-            HANDLE_RETURN(self._query_fn(
-                c_node, nodes_vec.data(), &count))
-        return [GraphNode._create(self._h_graph, nodes_vec[i])
+            HANDLE_RETURN(self._query_fn(c_node, nodes, &count))
+        return [GraphNode._create(self._h_graph, nodes[i])
                 for i in range(count)]
 
     cdef bint contains(self, GraphNode other):
         cdef cydriver.CUgraphNode c_node = as_cu(self._h_node)
         cdef cydriver.CUgraphNode target = as_cu(other._h_node)
         if c_node == NULL or target == NULL:
             return False
-        cdef cydriver.CUgraphNode buf[16]
-        cdef size_t count = 16
+        cdef cydriver.CUgraphNode stack_buf[16]
+        cdef cydriver.CUgraphNode* nodes
+        cdef size_t count = 0
         cdef size_t i
         with nogil:
-            HANDLE_RETURN(self._query_fn(c_node, buf, &count))
-
-        # Fast path for small sets.
-        if count <= 16:
-            for i in range(count):
-                if buf[i] == target:
-                    return True
+            HANDLE_RETURN(self._query_fn(c_node, NULL, &count))
+        if count == 0:
             return False
-
-        # Fallback for large sets.
         cdef vector[cydriver.CUgraphNode] nodes_vec
-        nodes_vec.resize(count)
+        if count <= 16:
+            nodes = stack_buf
+        else:
+            nodes_vec.resize(count)
+            nodes = nodes_vec.data()
         with nogil:
-            HANDLE_RETURN(self._query_fn(c_node, nodes_vec.data(), &count))
-        assert count == nodes_vec.size()
+            HANDLE_RETURN(self._query_fn(c_node, nodes, &count))
         for i in range(count):
-            if nodes_vec[i] == target:
+            if nodes[i] == target:
                 return True
         return False
 

cuda_core/cuda/core/graph/_graph_definition.pyx

@@ -361,19 +361,17 @@ cdef class GraphDefinition:
             All nodes in the graph.
         """
         cdef vector[cydriver.CUgraphNode] nodes_vec
-        nodes_vec.resize(128)
-        cdef size_t num_nodes = 128
+        cdef size_t num_nodes = 0
 
         with nogil:
-            HANDLE_RETURN(cydriver.cuGraphGetNodes(as_cu(self._h_graph), nodes_vec.data(), &num_nodes))
+            HANDLE_RETURN(cydriver.cuGraphGetNodes(as_cu(self._h_graph), NULL, &num_nodes))
 
         if num_nodes == 0:
             return set()
 
-        if num_nodes > 128:
-            nodes_vec.resize(num_nodes)
-            with nogil:
-                HANDLE_RETURN(cydriver.cuGraphGetNodes(as_cu(self._h_graph), nodes_vec.data(), &num_nodes))
+        nodes_vec.resize(num_nodes)
+        with nogil:
+            HANDLE_RETURN(cydriver.cuGraphGetNodes(as_cu(self._h_graph), nodes_vec.data(), &num_nodes))
 
         return {GraphNode._create(self._h_graph, nodes_vec[i]) for i in range(num_nodes)}
 
@@ -388,31 +386,28 @@ cdef class GraphDefinition:
         """
         cdef vector[cydriver.CUgraphNode] from_nodes
         cdef vector[cydriver.CUgraphNode] to_nodes
-        from_nodes.resize(128)
-        to_nodes.resize(128)
-        cdef size_t num_edges = 128
+        cdef size_t num_edges = 0
 
         with nogil:
             IF CUDA_CORE_BUILD_MAJOR >= 13:
                 HANDLE_RETURN(cydriver.cuGraphGetEdges(
-                    as_cu(self._h_graph), from_nodes.data(), to_nodes.data(), NULL, &num_edges))
+                    as_cu(self._h_graph), NULL, NULL, NULL, &num_edges))
             ELSE:
                 HANDLE_RETURN(cydriver.cuGraphGetEdges(
-                    as_cu(self._h_graph), from_nodes.data(), to_nodes.data(), &num_edges))
+                    as_cu(self._h_graph), NULL, NULL, &num_edges))
 
         if num_edges == 0:
             return set()
 
-        if num_edges > 128:
-            from_nodes.resize(num_edges)
-            to_nodes.resize(num_edges)
-            with nogil:
-                IF CUDA_CORE_BUILD_MAJOR >= 13:
-                    HANDLE_RETURN(cydriver.cuGraphGetEdges(
-                        as_cu(self._h_graph), from_nodes.data(), to_nodes.data(), NULL, &num_edges))
-                ELSE:
-                    HANDLE_RETURN(cydriver.cuGraphGetEdges(
-                        as_cu(self._h_graph), from_nodes.data(), to_nodes.data(), &num_edges))
+        from_nodes.resize(num_edges)
+        to_nodes.resize(num_edges)
+        with nogil:
+            IF CUDA_CORE_BUILD_MAJOR >= 13:
+                HANDLE_RETURN(cydriver.cuGraphGetEdges(
+                    as_cu(self._h_graph), from_nodes.data(), to_nodes.data(), NULL, &num_edges))
+            ELSE:
+                HANDLE_RETURN(cydriver.cuGraphGetEdges(
+                    as_cu(self._h_graph), from_nodes.data(), to_nodes.data(), &num_edges))
 
         return {
             (GraphNode._create(self._h_graph, from_nodes[i]),

cuda_core/tests/graph/test_graph_definition.py

@@ -633,6 +633,43 @@ def test_succ(nonempty_graph_spec):
         assert actual == spec.expected_succ[name], f"succ mismatch for node {name}"
 
 
+@pytest.mark.parametrize("adjacency_name", ("pred", "succ"))
+@pytest.mark.agent_authored(model="gpt-5.6")
+def test_large_adjacency_set_is_not_truncated(init_cuda, adjacency_name):
+    """Adjacency queries return and remove edges beyond the old 16-edge buffer."""
+    g = GraphDefinition()
+    hub = g.empty()
+    neighbors = [g.empty() for _ in range(20)]
+    adjacency = getattr(hub, adjacency_name)
+    adjacency.update(neighbors)
+
+    expected_edges = (
+        {(node, hub) for node in neighbors} if adjacency_name == "pred" else {(hub, node) for node in neighbors}
+    )
+    assert len(adjacency) == 20
+    assert set(adjacency) == set(neighbors)
+    assert neighbors[-1] in adjacency
+    assert g.edges() == expected_edges
+
+    adjacency.clear()
+    assert len(adjacency) == 0
+    assert g.edges() == set()
+
+
+@pytest.mark.agent_authored(model="gpt-5.6")
+def test_large_graph_queries_are_not_truncated(init_cuda):
+    """Graph queries return nodes and edges beyond the old 128-item buffers."""
+    g = GraphDefinition()
+    nodes = [g.empty() for _ in range(130)]
+    nodes[0].succ.update(nodes[1:])
+    nodes[1].succ.add(nodes[2])
+
+    expected_edges = {(nodes[0], node) for node in nodes[1:]}
+    expected_edges.add((nodes[1], nodes[2]))
+    assert g.nodes() == set(nodes)
+    assert g.edges() == expected_edges
+
+
 def test_node_graph_property(nonempty_graph_spec):
     """Every node's .graph property returns the parent GraphDefinition."""
     spec = nonempty_graph_spec