diff --git a/go/cstx.go b/go/cstx.go index e033ef8..4ef292a 100644 --- a/go/cstx.go +++ b/go/cstx.go @@ -33,10 +33,11 @@ type CSTX struct { eng engine projectID string - // Extensions, Graph, and Repo are lightweight namespaces sharing + // Extensions, Graph, Rag, and Repo are lightweight namespaces sharing // this runtime's state. Extensions *Extensions Graph *Graph + Rag *Rag Repo *Repository mu sync.Mutex @@ -60,6 +61,7 @@ func wrapRuntime(eng engine, projectID string) *CSTX { rt := &CSTX{eng: eng, projectID: projectID} rt.Extensions = &Extensions{eng: eng} rt.Graph = &Graph{eng: eng} + rt.Rag = &Rag{eng: eng} rt.Repo = &Repository{eng: eng} return rt } diff --git a/go/cstx_ffi.h b/go/cstx_ffi.h index ae52695..11fef2f 100644 --- a/go/cstx_ffi.h +++ b/go/cstx_ffi.h @@ -141,6 +141,21 @@ CstxStatusCode cstx_extension_anchor_concepts(struct CstxHandle *handle, struct CstxBuffer *output, struct CstxBuffer *error); +/** + * Parse one artifact through a native extension without mutating this graph. + * + * Input is a `ParserPayload`; output is a `Graph` batch plus the number of + * records parsed. The batch enters the same merge/link path as a parser + * implemented in any other language, so callers feed it to + * `cstx_graph_add_nodes` and then `cstx_graph_link`. The extension that owns + * the artifact must be enabled first or this reports `CSTX_NOT_FOUND`. + */ +CstxStatusCode cstx_graph_parse(struct CstxHandle *handle, + struct CstxSlice payload, + uint64_t *records, + struct CstxBuffer *output, + struct CstxBuffer *error); + /** * Add or merge a protobuf graph aggregate at the Rust-owned semantic boundary. */ diff --git a/go/engine.go b/go/engine.go index 729e75b..8b346d3 100644 --- a/go/engine.go +++ b/go/engine.go @@ -26,6 +26,7 @@ type engine interface { extensionParsesArtifact(context.Context, string) (bool, error) extensionAnchorConcepts(context.Context) (cstxproto.AnchorConceptCatalog, error) + graphParse(context.Context, *cstxproto.ParserPayload) (*cstxproto.Graph, uint64, error) graphAddNodes(context.Context, []*cstxproto.Node) (uint64, error) graphReplaceNodes(context.Context, []*cstxproto.Node) (uint64, error) graphAddRelationships(context.Context, []*cstxproto.Relationship) (uint64, error) @@ -33,17 +34,38 @@ type engine interface { graphDeleteNodes(context.Context, []string) (uint64, error) graphDeleteRelationships(context.Context, []string) (uint64, error) graphNode(context.Context, string) (*cstxproto.Node, error) + graphFindNode(context.Context, string) (*cstxproto.Node, error) graphRelationship(context.Context, string) (*cstxproto.Relationship, error) graphContains(context.Context, string) (bool, error) graphNodeCount(context.Context) (uint64, error) graphRelationshipCount(context.Context) (uint64, error) + graphNodeTypes(context.Context) ([]string, error) + graphDegree(context.Context, string, string) (uint64, error) graphStats(context.Context) (*cstxproto.GraphStats, error) graphNodes(context.Context, *cstxproto.NodeQuery) (graphCursor, error) graphRelationships(context.Context, *cstxproto.RelationshipQuery) (graphCursor, error) graphNeighbors(context.Context, *cstxproto.NeighborQuery) (graphCursor, error) graphQuery(context.Context, *cstxproto.GraphQuery) (graphCursor, error) graphAnalyze(context.Context, *cstxproto.Algorithm, *string) (uint8, bool, graphCursor, error) + graphLink(context.Context, []string, string) (*cstxproto.GraphLinkResult, error) + graphUpdateNodeFlags(context.Context, *cstxproto.NodeFlagChange) (uint64, error) + graphPatchNodeAnnotations(context.Context, *cstxproto.NodeAnnotationUpdate) (uint64, error) + graphFindAnchors(context.Context, string) (*cstxproto.GraphAnchorCatalog, error) + + // Derived-graph operations return an independently owned engine, matching + // the Python methods that return a new CSTX. graphSubgraph(context.Context, []string, uint32) (engine, error) + graphQuerySubgraph(context.Context, *cstxproto.GraphQuery) (engine, error) + graphInducedSubgraph(context.Context, []string, []string) (engine, error) + graphFilter(context.Context, *cstxproto.NodeFilter) (engine, error) + graphFilterWithReasons(context.Context, *cstxproto.NodeFilter) (engine, *cstxproto.GraphProjectionReport, error) + graphElevate(context.Context, string) (engine, error) + graphUnion(context.Context, engine) (engine, error) + graphDifference(context.Context, engine, string) (engine, error) + graphMerge(context.Context, engine) (uint64, error) + + ragIndex(context.Context, *cstxproto.RagIndexPlan) (ragIndexSession, error) + ragRetrieve(context.Context, *cstxproto.RagQuery) (ragRetrieval, error) repoResolve(context.Context, string) (string, error) repoHead(context.Context, string) (*string, error) @@ -70,3 +92,28 @@ type graphCursor interface { page(context.Context, int, int) (*cstxproto.GraphResultPage, error) close() } + +// ragIndexSession is a retained projection. It is a handle rather than a value +// because the records it holds are streamed in bounded pages instead of being +// materialized at once. +type ragIndexSession interface { + metadata(context.Context) (*cstxproto.RagIndexResult, error) + pending(context.Context, int, int) (*cstxproto.RagRecordPage, error) + deletes(context.Context) ([]string, error) + records(context.Context) (ragRecordIterator, error) + close() +} + +type ragRecordIterator interface { + next(context.Context) (*cstxproto.RagRecord, bool, error) + close() +} + +// ragRetrieval is a suspended retrieval: the plan is read with requests, the +// embedder's answers are handed back through complete, and completing consumes +// the retrieval. +type ragRetrieval interface { + requests(context.Context) (*cstxproto.RecallPlan, error) + complete(context.Context, *cstxproto.RecallResults) (*cstxproto.RagResult, error) + close() +} diff --git a/go/engine_native.go b/go/engine_native.go index ef18ae6..375c852 100644 --- a/go/engine_native.go +++ b/go/engine_native.go @@ -59,6 +59,26 @@ func (e *nativeEngine) close() error { return nil } +// adoptHandle takes ownership of a runtime handle the native side just +// produced. Every derived-graph operation returns one, so the finalizer is set +// in one place instead of at each call site. +func adoptHandle(handle *C.CstxHandle) engine { + derived := &nativeEngine{handle: handle} + runtime.SetFinalizer(derived, (*nativeEngine).finalize) + return derived +} + +// peerHandle reads the native handle out of another engine. The binary graph +// operations are the only place one runtime reaches into another, and a +// non-native peer cannot satisfy them. +func peerHandle(other engine, operation string) (*C.CstxHandle, error) { + native, ok := other.(*nativeEngine) + if !ok || native.handle == nil { + return nil, &Error{Code: CodeInvalidArgument, Operation: operation, Message: "other graph is not an open native runtime"} + } + return native.handle, nil +} + func (e *nativeEngine) graphSubgraph(_ context.Context, seedIDs []string, depth uint32) (engine, error) { payload, err := proto.Marshal(&cstxproto.GraphSelection{NodeIds: seedIDs}) if err != nil { @@ -73,9 +93,133 @@ func (e *nativeEngine) graphSubgraph(_ context.Context, seedIDs []string, depth if err != nil { return nil, err } - derived := &nativeEngine{handle: handle} - runtime.SetFinalizer(derived, (*nativeEngine).finalize) - return derived, nil + return adoptHandle(handle), nil +} + +// derivedHandle runs one native call that produces a new runtime handle from a +// serialized request. +func (e *nativeEngine) derivedHandle( + operation string, + request []byte, + call func(handle *C.CstxHandle, payload C.CstxSlice, output **C.CstxHandle, errBuf *C.CstxBuffer) C.CstxStatusCode, +) (engine, error) { + var handle *C.CstxHandle + err := statusCall(operation, func(errBuf *C.CstxBuffer) C.CstxStatusCode { + rc := call(e.handle, byteSlice(request), &handle, errBuf) + runtime.KeepAlive(request) + return rc + }) + if err != nil { + return nil, err + } + return adoptHandle(handle), nil +} + +func (e *nativeEngine) graphQuerySubgraph(_ context.Context, query *cstxproto.GraphQuery) (engine, error) { + payload, err := proto.Marshal(query) + if err != nil { + return nil, err + } + return e.derivedHandle("graph.query_subgraph", payload, + func(handle *C.CstxHandle, request C.CstxSlice, output **C.CstxHandle, errBuf *C.CstxBuffer) C.CstxStatusCode { + return C.cstx_graph_query_subgraph(handle, request, output, errBuf) + }) +} + +func (e *nativeEngine) graphInducedSubgraph(_ context.Context, nodeIDs, relationshipIDs []string) (engine, error) { + payload, err := proto.Marshal(&cstxproto.GraphSelection{NodeIds: nodeIDs, RelationshipIds: relationshipIDs}) + if err != nil { + return nil, err + } + return e.derivedHandle("graph.induced_subgraph", payload, + func(handle *C.CstxHandle, request C.CstxSlice, output **C.CstxHandle, errBuf *C.CstxBuffer) C.CstxStatusCode { + return C.cstx_graph_induced_subgraph(handle, request, output, errBuf) + }) +} + +func (e *nativeEngine) graphFilter(_ context.Context, filter *cstxproto.NodeFilter) (engine, error) { + payload, err := proto.Marshal(filter) + if err != nil { + return nil, err + } + return e.derivedHandle("graph.filter", payload, + func(handle *C.CstxHandle, request C.CstxSlice, output **C.CstxHandle, errBuf *C.CstxBuffer) C.CstxStatusCode { + return C.cstx_graph_filter(handle, request, output, errBuf) + }) +} + +func (e *nativeEngine) graphFilterWithReasons(_ context.Context, filter *cstxproto.NodeFilter) (engine, *cstxproto.GraphProjectionReport, error) { + payload, err := proto.Marshal(filter) + if err != nil { + return nil, nil, err + } + var handle *C.CstxHandle + details, err := bufferResult("graph.filter_with_reasons", func(out, errBuf *C.CstxBuffer) C.CstxStatusCode { + rc := C.cstx_graph_filter_with_reasons(e.handle, byteSlice(payload), &handle, out, errBuf) + runtime.KeepAlive(payload) + return rc + }) + if err != nil { + return nil, nil, err + } + var report cstxproto.GraphProjectionReport + if err := proto.Unmarshal(details, &report); err != nil { + return nil, nil, fmt.Errorf("cstx: decode graph projection report protobuf: %w", err) + } + return adoptHandle(handle), &report, nil +} + +func (e *nativeEngine) graphElevate(_ context.Context, conceptName string) (engine, error) { + var handle *C.CstxHandle + err := statusCall("graph.elevate", func(errBuf *C.CstxBuffer) C.CstxStatusCode { + rc := C.cstx_graph_elevate(e.handle, stringSlice(conceptName), &handle, errBuf) + runtime.KeepAlive(conceptName) + return rc + }) + if err != nil { + return nil, err + } + return adoptHandle(handle), nil +} + +func (e *nativeEngine) graphUnion(_ context.Context, other engine) (engine, error) { + right, err := peerHandle(other, "graph.union") + if err != nil { + return nil, err + } + var handle *C.CstxHandle + if err := statusCall("graph.union", func(errBuf *C.CstxBuffer) C.CstxStatusCode { + return C.cstx_graph_union(e.handle, right, &handle, errBuf) + }); err != nil { + return nil, err + } + return adoptHandle(handle), nil +} + +func (e *nativeEngine) graphDifference(_ context.Context, other engine, nodeType string) (engine, error) { + right, err := peerHandle(other, "graph.difference") + if err != nil { + return nil, err + } + var handle *C.CstxHandle + if err := statusCall("graph.difference", func(errBuf *C.CstxBuffer) C.CstxStatusCode { + rc := C.cstx_graph_difference(e.handle, right, optionalStringSlice(nodeType), &handle, errBuf) + runtime.KeepAlive(nodeType) + return rc + }); err != nil { + return nil, err + } + return adoptHandle(handle), nil +} + +func (e *nativeEngine) graphMerge(_ context.Context, other engine) (uint64, error) { + source, err := peerHandle(other, "graph.merge") + if err != nil { + return 0, err + } + return countResult("graph.merge", func(out *C.uint64_t, errBuf *C.CstxBuffer) C.CstxStatusCode { + return C.cstx_graph_merge(e.handle, source, out, errBuf) + }) } func (e *nativeEngine) graphDeleteNodes(_ context.Context, nodeIDs []string) (uint64, error) { @@ -370,6 +514,22 @@ func (e *nativeEngine) extensionAnchorConcepts(_ context.Context) (cstxproto.Anc // --- graph --------------------------------------------------------------- +func (e *nativeEngine) graphParse(_ context.Context, payload *cstxproto.ParserPayload) (*cstxproto.Graph, uint64, error) { + graph, records, err := e.graphParseWire(context.Background(), payload) + if err != nil { + return nil, 0, err + } + return &graph, records, nil +} + +func (e *nativeEngine) graphLink(_ context.Context, nodeIDs []string, dataSource string) (*cstxproto.GraphLinkResult, error) { + result, err := e.graphLinkWire(context.Background(), nodeIDs, dataSource) + if err != nil { + return nil, err + } + return &result, nil +} + func (e *nativeEngine) graphAddNodes(_ context.Context, nodes []*cstxproto.Node) (uint64, error) { return e.graphAddNodesWire(context.Background(), &cstxproto.Graph{Nodes: nodes}) } @@ -406,6 +566,85 @@ func (e *nativeEngine) graphRelationship(_ context.Context, relationshipID strin return &relationship, nil } +func (e *nativeEngine) graphFindNode(_ context.Context, identifier string) (*cstxproto.Node, error) { + data, err := bufferResult("graph.find_node", func(out, errBuf *C.CstxBuffer) C.CstxStatusCode { + rc := C.cstx_graph_find_node(e.handle, stringSlice(identifier), out, errBuf) + runtime.KeepAlive(identifier) + return rc + }) + if err != nil { + return nil, err + } + var node cstxproto.Node + if err := proto.Unmarshal(data, &node); err != nil { + return nil, fmt.Errorf("cstx: decode node protobuf: %w", err) + } + return &node, nil +} + +func (e *nativeEngine) graphNodeTypes(_ context.Context) ([]string, error) { + data, err := bufferResult("graph.node_types", func(out, errBuf *C.CstxBuffer) C.CstxStatusCode { + return C.cstx_graph_node_types(e.handle, out, errBuf) + }) + if err != nil { + return nil, err + } + var catalog cstxproto.NodeTypeCatalog + if err := proto.Unmarshal(data, &catalog); err != nil { + return nil, fmt.Errorf("cstx: decode node type catalog protobuf: %w", err) + } + return catalog.GetNodeTypes(), nil +} + +func (e *nativeEngine) graphDegree(_ context.Context, nodeID, direction string) (uint64, error) { + return countResult("graph.degree", func(out *C.uint64_t, errBuf *C.CstxBuffer) C.CstxStatusCode { + rc := C.cstx_graph_degree(e.handle, stringSlice(nodeID), stringSlice(direction), out, errBuf) + runtime.KeepAlive(nodeID) + runtime.KeepAlive(direction) + return rc + }) +} + +func (e *nativeEngine) graphUpdateNodeFlags(_ context.Context, change *cstxproto.NodeFlagChange) (uint64, error) { + payload, err := proto.Marshal(change) + if err != nil { + return 0, err + } + return countResult("graph.update_node_flags", func(out *C.uint64_t, errBuf *C.CstxBuffer) C.CstxStatusCode { + rc := C.cstx_graph_update_node_flags(e.handle, byteSlice(payload), out, errBuf) + runtime.KeepAlive(payload) + return rc + }) +} + +func (e *nativeEngine) graphPatchNodeAnnotations(_ context.Context, update *cstxproto.NodeAnnotationUpdate) (uint64, error) { + payload, err := proto.Marshal(update) + if err != nil { + return 0, err + } + return countResult("graph.patch_node_annotations", func(out *C.uint64_t, errBuf *C.CstxBuffer) C.CstxStatusCode { + rc := C.cstx_graph_patch_node_annotations(e.handle, byteSlice(payload), out, errBuf) + runtime.KeepAlive(payload) + return rc + }) +} + +func (e *nativeEngine) graphFindAnchors(_ context.Context, conceptName string) (*cstxproto.GraphAnchorCatalog, error) { + data, err := bufferResult("graph.find_anchors", func(out, errBuf *C.CstxBuffer) C.CstxStatusCode { + rc := C.cstx_graph_find_anchors(e.handle, stringSlice(conceptName), out, errBuf) + runtime.KeepAlive(conceptName) + return rc + }) + if err != nil { + return nil, err + } + var catalog cstxproto.GraphAnchorCatalog + if err := proto.Unmarshal(data, &catalog); err != nil { + return nil, fmt.Errorf("cstx: decode graph anchor catalog protobuf: %w", err) + } + return &catalog, nil +} + func (e *nativeEngine) graphContains(_ context.Context, nodeID string) (bool, error) { return boolResult("graph.contains", func(out *C.uint8_t, errBuf *C.CstxBuffer) C.CstxStatusCode { rc := C.cstx_graph_contains(e.handle, stringSlice(nodeID), out, errBuf) diff --git a/go/graph.go b/go/graph.go index aa98dda..59026a0 100644 --- a/go/graph.go +++ b/go/graph.go @@ -10,6 +10,24 @@ import ( // ingestion, and queries; the repository lifecycle lives elsewhere. type Graph struct{ eng engine } +// Parse runs one artifact through the native extension parser that owns it and +// returns a graph batch plus the number of records the parser read. The graph +// is not mutated: the batch enters the same merge/compute/link path as a parser +// implemented in any other language, so callers pass it to AddNodes and then +// Link. +// +// The extension providing the parser must be enabled first, or the call reports +// CodeNotFound. +func (g *Graph) Parse(ctx context.Context, payload *cstxproto.ParserPayload) (*cstxproto.Graph, uint64, error) { + if err := contextError(ctx); err != nil { + return nil, 0, err + } + if payload == nil { + return nil, 0, &Error{Code: CodeInvalidArgument, Operation: "graph.parse", Message: "payload must not be nil"} + } + return g.eng.graphParse(ctx, payload) +} + // AddNodes atomically adds or merges nodes and returns the number of elements // actually changed. A no-op write reports zero and does not invalidate // cursors. @@ -56,6 +74,17 @@ func (g *Graph) AddRelationship(ctx context.Context, relationship *cstxproto.Rel return g.eng.graphAddRelationship(ctx, relationship) } +// Link derives relationships between the selected nodes from the registered +// extension schemas and reports what it created or updated. dataSource names +// the origin recorded on the new relationships, so an unlinked batch stays +// traceable to the artifact it came from. +func (g *Graph) Link(ctx context.Context, nodeIDs []string, dataSource string) (*cstxproto.GraphLinkResult, error) { + if err := contextError(ctx); err != nil { + return nil, err + } + return g.eng.graphLink(ctx, nodeIDs, dataSource) +} + // DeleteNodes atomically removes nodes and all incident relationships. func (g *Graph) DeleteNodes(ctx context.Context, nodeIDs []string) (uint64, error) { if err := contextError(ctx); err != nil { @@ -80,6 +109,16 @@ func (g *Graph) Node(ctx context.Context, nodeID string) (*cstxproto.Node, error return g.eng.graphNode(ctx, nodeID) } +// FindNode resolves a human-readable identifier — an IP, a domain, a URL — +// to its node. Node takes a stable CSTX ID and FindNode takes the thing a +// user actually typed; both report CodeNotFound when nothing matches. +func (g *Graph) FindNode(ctx context.Context, identifier string) (*cstxproto.Node, error) { + if err := contextError(ctx); err != nil { + return nil, err + } + return g.eng.graphFindNode(ctx, identifier) +} + // Relationship returns one generated protobuf relationship or CodeNotFound. func (g *Graph) Relationship(ctx context.Context, relationshipID string) (*cstxproto.Relationship, error) { if err := contextError(ctx); err != nil { @@ -112,6 +151,62 @@ func (g *Graph) RelationshipCount(ctx context.Context) (uint64, error) { return g.eng.graphRelationshipCount(ctx) } +// NodeTypes returns the distinct node types currently present in the graph. +// It reports what the data holds, not what the registered extensions declare — +// use Extensions.Schemas for the declared catalog. +func (g *Graph) NodeTypes(ctx context.Context) ([]string, error) { + if err := contextError(ctx); err != nil { + return nil, err + } + return g.eng.graphNodeTypes(ctx) +} + +// Degree returns one node's relationship count in the given direction, which +// is "out", "in", or "both". An empty direction means "both". +func (g *Graph) Degree(ctx context.Context, nodeID, direction string) (uint64, error) { + if err := contextError(ctx); err != nil { + return 0, err + } + if direction == "" { + direction = "both" + } + return g.eng.graphDegree(ctx, nodeID, direction) +} + +// UpdateNodeFlags atomically applies one flag change to the selected nodes and +// returns how many changed. Flag bits are declared by extensions, not by this +// SDK; read them with FlagRegistry. +func (g *Graph) UpdateNodeFlags(ctx context.Context, change *cstxproto.NodeFlagChange) (uint64, error) { + if err := contextError(ctx); err != nil { + return 0, err + } + if change == nil { + return 0, &Error{Code: CodeInvalidArgument, Operation: "graph.update_node_flags", Message: "change must not be nil"} + } + return g.eng.graphUpdateNodeFlags(ctx, change) +} + +// PatchNodeAnnotations merges an annotation patch into the selected nodes and +// returns how many changed. An empty selection patches every node. +func (g *Graph) PatchNodeAnnotations(ctx context.Context, update *cstxproto.NodeAnnotationUpdate) (uint64, error) { + if err := contextError(ctx); err != nil { + return 0, err + } + if update == nil { + return 0, &Error{Code: CodeInvalidArgument, Operation: "graph.patch_node_annotations", Message: "update must not be nil"} + } + return g.eng.graphPatchNodeAnnotations(ctx, update) +} + +// FindAnchors returns the nodes anchoring a named concept. The concept names +// come from Extensions.AnchorConcepts. +func (g *Graph) FindAnchors(ctx context.Context, conceptName string) (*cstxproto.GraphAnchorCatalog, error) { + if err := contextError(ctx); err != nil { + return nil, err + } + return g.eng.graphFindAnchors(ctx, conceptName) +} + // Stats returns small aggregate counts. func (g *Graph) Stats(ctx context.Context) (*cstxproto.GraphStats, error) { if err := contextError(ctx); err != nil { @@ -220,3 +315,130 @@ func (g *Graph) Subgraph(ctx context.Context, seedIDs []string, depth uint32) (* } return wrapRuntime(eng, "derived"), nil } + +// QuerySubgraph runs a graph DSL query and materializes exactly the nodes and +// relationships it traversed as an independently owned runtime. Query returns +// a cursor over terminal nodes; this keeps the whole traced path. The caller +// must close the result. +func (g *Graph) QuerySubgraph(ctx context.Context, query *cstxproto.GraphQuery) (*CSTX, error) { + if err := contextError(ctx); err != nil { + return nil, err + } + if query == nil { + return nil, &Error{Code: CodeInvalidArgument, Operation: "graph.query_subgraph", Message: "query must not be nil"} + } + eng, err := g.eng.graphQuerySubgraph(ctx, query) + if err != nil { + return nil, err + } + return wrapRuntime(eng, "derived"), nil +} + +// InducedSubgraph materializes the named nodes as an independently owned +// runtime. A nil relationshipIDs keeps every relationship between those nodes; +// a non-nil one keeps only the relationships named. The caller must close the +// result. +func (g *Graph) InducedSubgraph(ctx context.Context, nodeIDs, relationshipIDs []string) (*CSTX, error) { + if err := contextError(ctx); err != nil { + return nil, err + } + eng, err := g.eng.graphInducedSubgraph(ctx, nodeIDs, relationshipIDs) + if err != nil { + return nil, err + } + return wrapRuntime(eng, "derived"), nil +} + +// Filter projects the graph through a node filter and returns the result as an +// independently owned runtime. The caller must close it. +func (g *Graph) Filter(ctx context.Context, filter *cstxproto.NodeFilter) (*CSTX, error) { + if err := contextError(ctx); err != nil { + return nil, err + } + if filter == nil { + filter = &cstxproto.NodeFilter{} + } + eng, err := g.eng.graphFilter(ctx, filter) + if err != nil { + return nil, err + } + return wrapRuntime(eng, "derived"), nil +} + +// FilterWithReasons is Filter plus a report naming every excluded node and why +// it was excluded, and whether the projection reused an existing one. Use it +// when a filtered result has to be explainable; use Filter when it does not. +// The caller must close the returned runtime. +func (g *Graph) FilterWithReasons(ctx context.Context, filter *cstxproto.NodeFilter) (*CSTX, *cstxproto.GraphProjectionReport, error) { + if err := contextError(ctx); err != nil { + return nil, nil, err + } + if filter == nil { + filter = &cstxproto.NodeFilter{} + } + eng, report, err := g.eng.graphFilterWithReasons(ctx, filter) + if err != nil { + return nil, nil, err + } + return wrapRuntime(eng, "derived"), report, nil +} + +// Elevate promotes a named concept's anchors into a graph organized around +// that concept and returns it as an independently owned runtime. The caller +// must close it. +func (g *Graph) Elevate(ctx context.Context, conceptName string) (*CSTX, error) { + if err := contextError(ctx); err != nil { + return nil, err + } + eng, err := g.eng.graphElevate(ctx, conceptName) + if err != nil { + return nil, err + } + return wrapRuntime(eng, "derived"), nil +} + +// Union returns a new runtime holding everything in this graph and in other. +// Neither input is modified; the caller must close the result. +func (g *Graph) Union(ctx context.Context, other *Graph) (*CSTX, error) { + if err := contextError(ctx); err != nil { + return nil, err + } + if other == nil { + return nil, &Error{Code: CodeInvalidArgument, Operation: "graph.union", Message: "other must not be nil"} + } + eng, err := g.eng.graphUnion(ctx, other.eng) + if err != nil { + return nil, err + } + return wrapRuntime(eng, "derived"), nil +} + +// Difference returns a new runtime holding the nodes present in this graph but +// not in other. A non-empty nodeType restricts the comparison to that type. +// Neither input is modified; the caller must close the result. +func (g *Graph) Difference(ctx context.Context, other *Graph, nodeType string) (*CSTX, error) { + if err := contextError(ctx); err != nil { + return nil, err + } + if other == nil { + return nil, &Error{Code: CodeInvalidArgument, Operation: "graph.difference", Message: "other must not be nil"} + } + eng, err := g.eng.graphDifference(ctx, other.eng, nodeType) + if err != nil { + return nil, err + } + return wrapRuntime(eng, "derived"), nil +} + +// Merge folds other into this graph in place and returns the number of +// elements changed. Union builds a third graph and leaves both inputs alone; +// Merge writes into this one. +func (g *Graph) Merge(ctx context.Context, other *Graph) (uint64, error) { + if err := contextError(ctx); err != nil { + return 0, err + } + if other == nil { + return 0, &Error{Code: CodeInvalidArgument, Operation: "graph.merge", Message: "other must not be nil"} + } + return g.eng.graphMerge(ctx, other.eng) +} diff --git a/go/graph_parse_test.go b/go/graph_parse_test.go new file mode 100644 index 0000000..37c7d56 --- /dev/null +++ b/go/graph_parse_test.go @@ -0,0 +1,87 @@ +package cstx + +import ( + "testing" + + "github.com/chainreactors/libcstx/go/proto/cstxproto" +) + +// The native accelerator is only half of an ingest: Parse returns a batch the +// caller feeds through the same AddNodes/Link path a parser in any language +// uses. This walks that whole path, because either half alone passing does not +// show the batch is actually writable. +func TestParseAddNodesLink(t *testing.T) { + rt := openRuntime(t) + ctx := testContext + + data := []byte(`{"ip":"127.0.0.1","port":"443","protocol":"https","host":"example.com","uri":"/","title":"hello"}` + "\n") + batch, records, err := rt.Graph.Parse(ctx, &cstxproto.ParserPayload{ + Plugin: "easm", + Artifact: "gogo", + Data: data, + }) + if err != nil { + t.Fatalf("parse: %v", err) + } + if records != 1 { + t.Fatalf("records = %d; want 1", records) + } + // One gogo record yields an ip, a cidr, a port and an app. + if len(batch.Nodes) != 4 { + t.Fatalf("parsed %d nodes; want 4", len(batch.Nodes)) + } + + affected, err := rt.Graph.AddNodes(ctx, batch.Nodes) + if err != nil { + t.Fatalf("add nodes: %v", err) + } + if affected != 4 { + t.Fatalf("affected = %d; want 4", affected) + } + + ids := make([]string, 0, len(batch.Nodes)) + for _, node := range batch.Nodes { + ids = append(ids, node.GetId()) + } + link, err := rt.Graph.Link(ctx, ids, "graph_parse_test") + if err != nil { + t.Fatalf("link: %v", err) + } + if len(link.RelationshipIds) == 0 { + t.Fatalf("link created no relationships: %+v", link) + } + + node, err := rt.Graph.Node(ctx, ids[0]) + if err != nil { + t.Fatalf("node: %v", err) + } + if node.GetValue().GetNodeType() != "ip" { + t.Fatalf("node type = %q; want ip", node.GetValue().GetNodeType()) + } +} + +// Parse is the accelerator, not the write: an unenabled extension has no +// registered parser to dispatch to, and the batch is not committed either way. +func TestParseWithoutEnabledExtension(t *testing.T) { + rt, err := Open(testContext, &cstxproto.RuntimeConfig{ProjectId: "sdk-go-parse-disabled"}) + if err != nil { + t.Fatalf("open: %v", err) + } + t.Cleanup(func() { _ = rt.Close() }) + + _, _, err = rt.Graph.Parse(testContext, &cstxproto.ParserPayload{ + Plugin: "easm", + Artifact: "gogo", + Data: []byte(`{"ip":"127.0.0.1","port":"443"}` + "\n"), + }) + if err == nil { + t.Fatal("parse without an enabled extension reported success") + } + count, err := rt.Graph.NodeCount(testContext) + if err != nil { + t.Fatalf("node count: %v", err) + } + if count != 0 { + t.Fatalf("node count = %d; want 0", count) + } +} diff --git a/go/graph_surface_test.go b/go/graph_surface_test.go new file mode 100644 index 0000000..cc61e7b --- /dev/null +++ b/go/graph_surface_test.go @@ -0,0 +1,362 @@ +package cstx + +import ( + "errors" + "slices" + "testing" + + "github.com/chainreactors/libcstx/go/proto/cstxproto" + "google.golang.org/protobuf/types/known/structpb" +) + +// These cover the operations the Go SDK declared through cstx_ffi.h but never +// bound, so each one asserts the Rust semantics rather than only that the call +// returns. + +func TestFindNodeResolvesIdentifier(t *testing.T) { + rt := openRuntime(t) + addDomain(t, rt, "findable.example") + + node, err := rt.Graph.FindNode(testContext, "findable.example") + if err != nil { + t.Fatalf("find node: %v", err) + } + if got := domainValue(t, node); got != "findable.example" { + t.Fatalf("found %q; want findable.example", got) + } + + if _, err := rt.Graph.FindNode(testContext, "absent.example"); err == nil { + t.Fatal("find node resolved an identifier that is not in the graph") + } +} + +func TestNodeTypesReportsPresentTypesOnly(t *testing.T) { + rt := openRuntime(t) + + empty, err := rt.Graph.NodeTypes(testContext) + if err != nil { + t.Fatalf("node types: %v", err) + } + // easm is enabled, so the schema declares many types. None is present yet. + if len(empty) != 0 { + t.Fatalf("empty graph reported node types %v", empty) + } + + addDomain(t, rt, "typed.example") + types, err := rt.Graph.NodeTypes(testContext) + if err != nil { + t.Fatalf("node types: %v", err) + } + if !slices.Contains(types, "domain") { + t.Fatalf("node types = %v; want it to contain domain", types) + } +} + +func TestDegreeCountsByDirection(t *testing.T) { + rt := openRuntime(t) + addDomain(t, rt, "source.example") + addDomain(t, rt, "target.example") + if _, err := rt.Graph.AddRelationships(testContext, []*cstxproto.Relationship{ + usesRelationship("domain:source.example", "domain:target.example"), + }); err != nil { + t.Fatalf("add relationship: %v", err) + } + + for _, testCase := range []struct { + direction string + want uint64 + }{ + {"out", 1}, + {"in", 0}, + {"both", 1}, + // An empty direction is the documented "both". + {"", 1}, + } { + got, err := rt.Graph.Degree(testContext, "domain:source.example", testCase.direction) + if err != nil { + t.Fatalf("degree %q: %v", testCase.direction, err) + } + if got != testCase.want { + t.Fatalf("degree %q = %d; want %d", testCase.direction, got, testCase.want) + } + } +} + +func TestUpdateNodeFlagsMergesAndReplaces(t *testing.T) { + rt := openRuntime(t) + addDomain(t, rt, "flagged.example") + selection := &cstxproto.GraphSelection{NodeIds: []string{"domain:flagged.example"}} + + affected, err := rt.Graph.UpdateNodeFlags(testContext, &cstxproto.NodeFlagChange{ + Selection: selection, + Update: &cstxproto.NodeFlagUpdate{ + Mode: cstxproto.NodeFlagUpdateMode_NODE_FLAG_UPDATE_MERGE, + AddMask: 0b101, + }, + }) + if err != nil { + t.Fatalf("merge flags: %v", err) + } + if affected != 1 { + t.Fatalf("merge affected = %d; want 1", affected) + } + node, err := rt.Graph.Node(testContext, "domain:flagged.example") + if err != nil { + t.Fatalf("node: %v", err) + } + if node.GetFlagsMask() != 0b101 { + t.Fatalf("flags after merge = %b; want 101", node.GetFlagsMask()) + } + + // Replace overwrites the mask instead of OR-ing into it, which is the + // distinction the mode field exists for. + if _, err := rt.Graph.UpdateNodeFlags(testContext, &cstxproto.NodeFlagChange{ + Selection: selection, + Update: &cstxproto.NodeFlagUpdate{ + Mode: cstxproto.NodeFlagUpdateMode_NODE_FLAG_UPDATE_REPLACE, + ReplaceMask: 0b010, + }, + }); err != nil { + t.Fatalf("replace flags: %v", err) + } + node, err = rt.Graph.Node(testContext, "domain:flagged.example") + if err != nil { + t.Fatalf("node: %v", err) + } + if node.GetFlagsMask() != 0b010 { + t.Fatalf("flags after replace = %b; want 010", node.GetFlagsMask()) + } +} + +func TestPatchNodeAnnotationsMergesIntoSelection(t *testing.T) { + rt := openRuntime(t) + addDomain(t, rt, "annotated.example") + + affected, err := rt.Graph.PatchNodeAnnotations(testContext, &cstxproto.NodeAnnotationUpdate{ + Selection: &cstxproto.GraphSelection{NodeIds: []string{"domain:annotated.example"}}, + Annotations: &structpb.Struct{Fields: map[string]*structpb.Value{"owner": structpb.NewStringValue("go-test")}}, + }) + if err != nil { + t.Fatalf("patch annotations: %v", err) + } + if affected != 1 { + t.Fatalf("affected = %d; want 1", affected) + } + node, err := rt.Graph.Node(testContext, "domain:annotated.example") + if err != nil { + t.Fatalf("node: %v", err) + } + if got := node.GetAnnotations().GetFields()["owner"].GetStringValue(); got != "go-test" { + t.Fatalf("annotation owner = %q; want go-test", got) + } +} + +func TestUnionAndDifferenceLeaveInputsAlone(t *testing.T) { + left := openRuntime(t) + addDomain(t, left, "shared.example") + addDomain(t, left, "left-only.example") + + right := openRuntime(t) + addDomain(t, right, "shared.example") + + union, err := left.Graph.Union(testContext, right.Graph) + if err != nil { + t.Fatalf("union: %v", err) + } + t.Cleanup(func() { _ = union.Close() }) + count, err := union.Graph.NodeCount(testContext) + if err != nil { + t.Fatalf("union node count: %v", err) + } + if count != 2 { + t.Fatalf("union node count = %d; want 2", count) + } + + difference, err := left.Graph.Difference(testContext, right.Graph, "") + if err != nil { + t.Fatalf("difference: %v", err) + } + t.Cleanup(func() { _ = difference.Close() }) + if _, err := difference.Graph.Node(testContext, "domain:left-only.example"); err != nil { + t.Fatalf("difference is missing the left-only node: %v", err) + } + if contains, err := difference.Graph.Contains(testContext, "domain:shared.example"); err != nil { + t.Fatalf("difference contains: %v", err) + } else if contains { + t.Fatal("difference kept a node present in both graphs") + } + + // Neither input was modified by either derived graph. + leftCount, err := left.Graph.NodeCount(testContext) + if err != nil { + t.Fatalf("left node count: %v", err) + } + if leftCount != 2 { + t.Fatalf("left node count = %d; want 2", leftCount) + } +} + +func TestMergeWritesIntoTheTargetGraph(t *testing.T) { + target := openRuntime(t) + addDomain(t, target, "target-only.example") + + source := openRuntime(t) + addDomain(t, source, "merged-in.example") + + affected, err := target.Graph.Merge(testContext, source.Graph) + if err != nil { + t.Fatalf("merge: %v", err) + } + if affected == 0 { + t.Fatal("merge reported no change") + } + if contains, err := target.Graph.Contains(testContext, "domain:merged-in.example"); err != nil { + t.Fatalf("contains: %v", err) + } else if !contains { + t.Fatal("merge did not write the source node into the target") + } + // Unlike Union, the source is untouched and the target grew in place. + sourceCount, err := source.Graph.NodeCount(testContext) + if err != nil { + t.Fatalf("source node count: %v", err) + } + if sourceCount != 1 { + t.Fatalf("source node count = %d; want 1", sourceCount) + } +} + +func TestInducedSubgraphKeepsSelectedNodes(t *testing.T) { + rt := openRuntime(t) + addDomain(t, rt, "kept.example") + addDomain(t, rt, "dropped.example") + + induced, err := rt.Graph.InducedSubgraph(testContext, []string{"domain:kept.example"}, nil) + if err != nil { + t.Fatalf("induced subgraph: %v", err) + } + t.Cleanup(func() { _ = induced.Close() }) + + count, err := induced.Graph.NodeCount(testContext) + if err != nil { + t.Fatalf("induced node count: %v", err) + } + if count != 1 { + t.Fatalf("induced node count = %d; want 1", count) + } + if contains, err := induced.Graph.Contains(testContext, "domain:dropped.example"); err != nil { + t.Fatalf("contains: %v", err) + } else if contains { + t.Fatal("induced subgraph kept an unselected node") + } +} + +func TestFilterWithReasonsExplainsExclusions(t *testing.T) { + rt := openRuntime(t) + addDomain(t, rt, "included.example") + addDomain(t, rt, "excluded.example") + if _, err := rt.Graph.UpdateNodeFlags(testContext, &cstxproto.NodeFlagChange{ + Selection: &cstxproto.GraphSelection{NodeIds: []string{"domain:excluded.example"}}, + Update: &cstxproto.NodeFlagUpdate{ + Mode: cstxproto.NodeFlagUpdateMode_NODE_FLAG_UPDATE_MERGE, + AddMask: 0b1, + }, + }); err != nil { + t.Fatalf("flag node: %v", err) + } + + filter := &cstxproto.NodeFilter{FlagsNoneMask: 0b1} + filtered, err := rt.Graph.Filter(testContext, filter) + if err != nil { + t.Fatalf("filter: %v", err) + } + t.Cleanup(func() { _ = filtered.Close() }) + count, err := filtered.Graph.NodeCount(testContext) + if err != nil { + t.Fatalf("filtered node count: %v", err) + } + if count != 1 { + t.Fatalf("filtered node count = %d; want 1", count) + } + + // The reasons variant returns the same projection plus why each node went. + explained, report, err := rt.Graph.FilterWithReasons(testContext, filter) + if err != nil { + t.Fatalf("filter with reasons: %v", err) + } + t.Cleanup(func() { _ = explained.Close() }) + if len(report.GetExcludedNodes()) != 1 { + t.Fatalf("excluded %d node(s); want 1: %+v", len(report.GetExcludedNodes()), report) + } + excluded := report.GetExcludedNodes()[0] + if excluded.GetNodeId() != "domain:excluded.example" { + t.Fatalf("excluded node = %q; want domain:excluded.example", excluded.GetNodeId()) + } + if excluded.GetReason() == "" { + t.Fatal("exclusion carries no reason, which is the whole point of this variant") + } +} + +func TestQuerySubgraphMaterializesTracedPath(t *testing.T) { + rt := openRuntime(t) + addDomain(t, rt, "traced.example") + + subgraph, err := rt.Graph.QuerySubgraph(testContext, &cstxproto.GraphQuery{Expression: "domain"}) + if err != nil { + t.Fatalf("query subgraph: %v", err) + } + t.Cleanup(func() { _ = subgraph.Close() }) + + count, err := subgraph.Graph.NodeCount(testContext) + if err != nil { + t.Fatalf("subgraph node count: %v", err) + } + if count != 1 { + t.Fatalf("subgraph node count = %d; want 1", count) + } +} + +func TestFindAnchorsAndElevateShareConceptNames(t *testing.T) { + rt := openRuntime(t) + addDomain(t, rt, "anchor.example") + + concepts, err := rt.Extensions.AnchorConcepts(testContext) + if err != nil { + t.Fatalf("anchor concepts: %v", err) + } + if len(concepts.GetConcepts()) == 0 { + t.Skip("the enabled extensions declare no anchor concept") + } + concept := concepts.GetConcepts()[0].GetName() + + if _, err := rt.Graph.FindAnchors(testContext, concept); err != nil { + t.Fatalf("find anchors %q: %v", concept, err) + } + elevated, err := rt.Graph.Elevate(testContext, concept) + if err != nil { + t.Fatalf("elevate %q: %v", concept, err) + } + t.Cleanup(func() { _ = elevated.Close() }) + + // An unknown concept is an error, not an empty result. + if _, err := rt.Graph.FindAnchors(testContext, "not-a-concept"); err == nil { + t.Fatal("find anchors accepted an unknown concept") + } +} + +func TestBinaryGraphOperationsRejectNilPeer(t *testing.T) { + rt := openRuntime(t) + for name, call := range map[string]func() error{ + "union": func() error { _, err := rt.Graph.Union(testContext, nil); return err }, + "difference": func() error { _, err := rt.Graph.Difference(testContext, nil, ""); return err }, + "merge": func() error { _, err := rt.Graph.Merge(testContext, nil); return err }, + } { + err := call() + if err == nil { + t.Fatalf("%s accepted a nil peer", name) + } + var cstxErr *Error + if !errors.As(err, &cstxErr) || cstxErr.Code != CodeInvalidArgument { + t.Fatalf("%s error = %v; want CodeInvalidArgument", name, err) + } + } +} diff --git a/go/lib/darwin_amd64/libcstx_ffi.a b/go/lib/darwin_amd64/libcstx_ffi.a index 728b112..392c565 100644 Binary files a/go/lib/darwin_amd64/libcstx_ffi.a and b/go/lib/darwin_amd64/libcstx_ffi.a differ diff --git a/go/lib/darwin_arm64/libcstx_ffi.a b/go/lib/darwin_arm64/libcstx_ffi.a index ba9e7cc..c3c5d26 100644 Binary files a/go/lib/darwin_arm64/libcstx_ffi.a and b/go/lib/darwin_arm64/libcstx_ffi.a differ diff --git a/go/lib/linux_amd64/libcstx_ffi.a b/go/lib/linux_amd64/libcstx_ffi.a index 568f8bf..0f2b881 100644 Binary files a/go/lib/linux_amd64/libcstx_ffi.a and b/go/lib/linux_amd64/libcstx_ffi.a differ diff --git a/go/lib/linux_arm64/libcstx_ffi.a b/go/lib/linux_arm64/libcstx_ffi.a index cadec22..f28a3b4 100644 Binary files a/go/lib/linux_arm64/libcstx_ffi.a and b/go/lib/linux_arm64/libcstx_ffi.a differ diff --git a/go/lib/windows_amd64/libcstx_ffi.a b/go/lib/windows_amd64/libcstx_ffi.a index 75b82ec..de07f25 100644 Binary files a/go/lib/windows_amd64/libcstx_ffi.a and b/go/lib/windows_amd64/libcstx_ffi.a differ diff --git a/go/proto_native.go b/go/proto_native.go index 8de1d73..3027eed 100644 --- a/go/proto_native.go +++ b/go/proto_native.go @@ -17,6 +17,50 @@ import ( "google.golang.org/protobuf/proto" ) +func (e *nativeEngine) graphParseWire(_ context.Context, payload *cstxproto.ParserPayload) (cstxproto.Graph, uint64, error) { + if payload == nil { + return cstxproto.Graph{}, 0, &Error{Code: CodeInvalidArgument, Operation: "graph.parse", Message: "payload must not be nil"} + } + encoded, err := proto.Marshal(payload) + if err != nil { + return cstxproto.Graph{}, 0, err + } + var records C.uint64_t + var out, errBuf C.CstxBuffer + if err := statusError(C.cstx_graph_parse(e.handle, byteSlice(encoded), &records, &out, &errBuf), "graph.parse", &errBuf); err != nil { + C.cstx_buffer_free(&out) + return cstxproto.Graph{}, 0, err + } + runtime.KeepAlive(encoded) + data := takeBuffer(&out) + var graph cstxproto.Graph + if err := proto.Unmarshal(data, &graph); err != nil { + return cstxproto.Graph{}, 0, fmt.Errorf("cstx: decode graph protobuf: %w", err) + } + return graph, uint64(records), nil +} + +func (e *nativeEngine) graphLinkWire(_ context.Context, nodeIDs []string, dataSource string) (cstxproto.GraphLinkResult, error) { + selection, err := proto.Marshal(&cstxproto.GraphSelection{NodeIds: nodeIDs}) + if err != nil { + return cstxproto.GraphLinkResult{}, err + } + data, err := bufferResult("graph.link", func(out, errBuf *C.CstxBuffer) C.CstxStatusCode { + rc := C.cstx_graph_link(e.handle, byteSlice(selection), stringSlice(dataSource), out, errBuf) + runtime.KeepAlive(selection) + runtime.KeepAlive(dataSource) + return rc + }) + if err != nil { + return cstxproto.GraphLinkResult{}, err + } + var result cstxproto.GraphLinkResult + if err := proto.Unmarshal(data, &result); err != nil { + return cstxproto.GraphLinkResult{}, fmt.Errorf("cstx: decode graph link result protobuf: %w", err) + } + return result, nil +} + func (e *nativeEngine) graphAddNodesWire(_ context.Context, graph *cstxproto.Graph) (uint64, error) { payload, err := proto.Marshal(graph) if err != nil { diff --git a/go/rag.go b/go/rag.go new file mode 100644 index 0000000..6e2eb8b --- /dev/null +++ b/go/rag.go @@ -0,0 +1,216 @@ +package cstx + +import ( + "context" + + "github.com/chainreactors/libcstx/go/proto/cstxproto" +) + +// Rag is the retrieval namespace of a CSTX runtime. It owns two suspended +// operations rather than two calls: Index projects the graph into records the +// caller embeds and stores, and Retrieve plans recall the caller executes +// against its own vector store. CSTX never talks to an embedder or an index; +// it decides what to ask for and what to do with the answers. +type Rag struct{ eng engine } + +// Index projects the graph into a retained session of records to upsert and +// IDs to delete. Nothing is written anywhere: the caller streams the records +// out, embeds and stores them, and closes the session. +func (r *Rag) Index(ctx context.Context, plan *cstxproto.RagIndexPlan) (*RagIndexSession, error) { + if err := contextError(ctx); err != nil { + return nil, err + } + if plan == nil { + return nil, &Error{Code: CodeInvalidArgument, Operation: "graph.rag.index", Message: "plan must not be nil"} + } + session, err := r.eng.ragIndex(ctx, plan) + if err != nil { + return nil, err + } + return &RagIndexSession{inner: session}, nil +} + +// Retrieve suspends a retrieval against one graph generation and checkpoint. +// Read the plan with Requests, run it against the vector store, and hand the +// answers back through Complete. +func (r *Rag) Retrieve(ctx context.Context, query *cstxproto.RagQuery) (*RagRetrieval, error) { + if err := contextError(ctx); err != nil { + return nil, err + } + if query == nil { + return nil, &Error{Code: CodeInvalidArgument, Operation: "graph.rag.retrieve", Message: "query must not be nil"} + } + retrieval, err := r.eng.ragRetrieve(ctx, query) + if err != nil { + return nil, err + } + return &RagRetrieval{inner: retrieval}, nil +} + +// RagIndexSession is a retained deterministic projection. It holds native +// state until closed, so callers should Close it when done; repeated calls are +// safe. +type RagIndexSession struct { + inner ragIndexSession + done bool +} + +func (s *RagIndexSession) closedError(operation string) error { + return &Error{Code: CodeNotInitialized, Operation: operation, Message: "index session is closed"} +} + +// Metadata returns the projection's identity and counts: the idempotent +// operation ID, the commit it targets, whether it replaces the whole index, +// and how many records it upserts and deletes. +func (s *RagIndexSession) Metadata(ctx context.Context) (*cstxproto.RagIndexResult, error) { + if s.done { + return nil, s.closedError("graph.rag.index.metadata") + } + if err := contextError(ctx); err != nil { + return nil, err + } + return s.inner.metadata(ctx) +} + +// Pending returns one bounded page of projected records. Records streams the +// same records without paging arithmetic; Pending is for a caller that wants +// to control the window itself. +func (s *RagIndexSession) Pending(ctx context.Context, offset, limit int) (*cstxproto.RagRecordPage, error) { + if s.done { + return nil, s.closedError("graph.rag.index.pending") + } + if err := contextError(ctx); err != nil { + return nil, err + } + if offset < 0 || limit < 0 { + return nil, &Error{Code: CodeInvalidArgument, Operation: "graph.rag.index.pending", Message: "offset and limit must not be negative"} + } + return s.inner.pending(ctx, offset, limit) +} + +// Deletes returns the record IDs this projection removes from the index. +func (s *RagIndexSession) Deletes(ctx context.Context) ([]string, error) { + if s.done { + return nil, s.closedError("graph.rag.index.deletes") + } + if err := contextError(ctx); err != nil { + return nil, err + } + return s.inner.deletes(ctx) +} + +// Records streams every projected record through a bounded native iterator, +// which never materializes the whole projection in memory. The caller must +// close the iterator. +func (s *RagIndexSession) Records(ctx context.Context) (*RagRecordIterator, error) { + if s.done { + return nil, s.closedError("graph.rag.index.records") + } + if err := contextError(ctx); err != nil { + return nil, err + } + iterator, err := s.inner.records(ctx) + if err != nil { + return nil, err + } + return &RagRecordIterator{inner: iterator}, nil +} + +// Close releases the retained projection; repeated calls are safe. +func (s *RagIndexSession) Close() error { + if !s.done { + s.done = true + s.inner.close() + } + return nil +} + +// Closed reports whether Close has been called. +func (s *RagIndexSession) Closed() bool { return s.done } + +// RagRecordIterator walks the projected records of one index session. +type RagRecordIterator struct { + inner ragRecordIterator + done bool +} + +// Next returns the next record. The boolean is false at the end of the +// projection, which is not an error. +func (i *RagRecordIterator) Next(ctx context.Context) (*cstxproto.RagRecord, bool, error) { + if i.done { + return nil, false, &Error{Code: CodeInvalidArgument, Operation: "graph.rag.index.records.next", Message: "iterator is closed"} + } + if err := contextError(ctx); err != nil { + return nil, false, err + } + return i.inner.next(ctx) +} + +// Close releases the iterator early; repeated calls are safe. +func (i *RagRecordIterator) Close() error { + if !i.done { + i.done = true + i.inner.close() + } + return nil +} + +// Closed reports whether Close has been called. +func (i *RagRecordIterator) Closed() bool { return i.done } + +// RagRetrieval is a suspended retrieval bound to one graph generation and +// checkpoint. Complete consumes it. +type RagRetrieval struct { + inner ragRetrieval + done bool + completed bool +} + +// Requests returns the recall plan: what the caller must look up in its own +// vector store before CSTX can finish the retrieval. +func (r *RagRetrieval) Requests(ctx context.Context) (*cstxproto.RecallPlan, error) { + if r.done || r.completed { + return nil, &Error{Code: CodeNotInitialized, Operation: "graph.rag.retrieve.requests", Message: "retrieval is closed or completed"} + } + if err := contextError(ctx); err != nil { + return nil, err + } + return r.inner.requests(ctx) +} + +// Complete fuses the recall results into a final answer and consumes the +// retrieval. A nil results means the caller ran no external recall, which is +// the normal case when the built-in lexical projection is the only source — +// that recall never leaves Rust. Calling Complete twice reports CodeConflict. +func (r *RagRetrieval) Complete(ctx context.Context, results *cstxproto.RecallResults) (*cstxproto.RagResult, error) { + if r.done { + return nil, &Error{Code: CodeNotInitialized, Operation: "graph.rag.complete", Message: "retrieval is closed"} + } + if r.completed { + return nil, &Error{Code: CodeConflict, Operation: "graph.rag.complete", Message: "retrieval has already completed"} + } + if err := contextError(ctx); err != nil { + return nil, err + } + if results == nil { + results = &cstxproto.RecallResults{} + } + result, err := r.inner.complete(ctx, results) + if err != nil { + return nil, err + } + r.completed = true + return result, nil +} + +// Close releases the suspended retrieval; repeated calls are safe. +func (r *RagRetrieval) Close() error { + if !r.done { + r.done = true + r.inner.close() + } + return nil +} + +// Closed reports whether Close has been called. +func (r *RagRetrieval) Closed() bool { return r.done } diff --git a/go/rag_native.go b/go/rag_native.go new file mode 100644 index 0000000..fca4e22 --- /dev/null +++ b/go/rag_native.go @@ -0,0 +1,244 @@ +package cstx + +// RAG crosses the ABI as three retained native objects rather than as plain +// calls: an index session holding a projection, a bounded iterator over that +// projection's records, and a suspended retrieval. Each owns a C pointer, so +// each gets the same explicit-close-plus-finalizer treatment as a graph cursor. + +/* +#include "cstx_ffi.h" +*/ +import "C" + +import ( + "context" + "fmt" + "runtime" + + "github.com/chainreactors/libcstx/go/proto/cstxproto" + "google.golang.org/protobuf/proto" +) + +func (e *nativeEngine) ragIndex(_ context.Context, plan *cstxproto.RagIndexPlan) (ragIndexSession, error) { + payload, err := proto.Marshal(plan) + if err != nil { + return nil, err + } + var session *C.CstxRagIndexSession + err = statusCall("graph.rag.index", func(errBuf *C.CstxBuffer) C.CstxStatusCode { + rc := C.cstx_rag_index(e.handle, byteSlice(payload), &session, errBuf) + runtime.KeepAlive(payload) + return rc + }) + if err != nil { + return nil, err + } + return newNativeRagIndexSession(session), nil +} + +func (e *nativeEngine) ragRetrieve(_ context.Context, query *cstxproto.RagQuery) (ragRetrieval, error) { + payload, err := proto.Marshal(query) + if err != nil { + return nil, err + } + var retrieval *C.CstxRagRetrieval + err = statusCall("graph.rag.retrieve", func(errBuf *C.CstxBuffer) C.CstxStatusCode { + rc := C.cstx_rag_retrieve(e.handle, byteSlice(payload), &retrieval, errBuf) + runtime.KeepAlive(payload) + return rc + }) + if err != nil { + return nil, err + } + return newNativeRagRetrieval(retrieval), nil +} + +// --- index session ------------------------------------------------------- + +type nativeRagIndexSession struct{ session *C.CstxRagIndexSession } + +func newNativeRagIndexSession(session *C.CstxRagIndexSession) *nativeRagIndexSession { + result := &nativeRagIndexSession{session: session} + runtime.SetFinalizer(result, (*nativeRagIndexSession).close) + return result +} + +func (s *nativeRagIndexSession) closedError(operation string) error { + return &Error{Code: CodeNotInitialized, Operation: operation, Message: "index session is closed"} +} + +func (s *nativeRagIndexSession) metadata(_ context.Context) (*cstxproto.RagIndexResult, error) { + if s.session == nil { + return nil, s.closedError("graph.rag.index.metadata") + } + data, err := bufferResult("graph.rag.index.metadata", func(out, errBuf *C.CstxBuffer) C.CstxStatusCode { + return C.cstx_rag_index_session_metadata(s.session, out, errBuf) + }) + if err != nil { + return nil, err + } + var result cstxproto.RagIndexResult + if err := proto.Unmarshal(data, &result); err != nil { + return nil, fmt.Errorf("cstx: decode rag index result protobuf: %w", err) + } + return &result, nil +} + +func (s *nativeRagIndexSession) pending(_ context.Context, offset, limit int) (*cstxproto.RagRecordPage, error) { + if s.session == nil { + return nil, s.closedError("graph.rag.index.pending") + } + data, err := bufferResult("graph.rag.index.pending", func(out, errBuf *C.CstxBuffer) C.CstxStatusCode { + return C.cstx_rag_index_session_pending(s.session, C.size_t(offset), C.size_t(limit), out, errBuf) + }) + if err != nil { + return nil, err + } + var page cstxproto.RagRecordPage + if err := proto.Unmarshal(data, &page); err != nil { + return nil, fmt.Errorf("cstx: decode rag record page protobuf: %w", err) + } + return &page, nil +} + +func (s *nativeRagIndexSession) deletes(_ context.Context) ([]string, error) { + if s.session == nil { + return nil, s.closedError("graph.rag.index.deletes") + } + data, err := bufferResult("graph.rag.index.deletes", func(out, errBuf *C.CstxBuffer) C.CstxStatusCode { + return C.cstx_rag_index_session_deletes(s.session, out, errBuf) + }) + if err != nil { + return nil, err + } + var selection cstxproto.ObjectSelection + if err := proto.Unmarshal(data, &selection); err != nil { + return nil, fmt.Errorf("cstx: decode object selection protobuf: %w", err) + } + return selection.GetObjectIds(), nil +} + +func (s *nativeRagIndexSession) records(_ context.Context) (ragRecordIterator, error) { + if s.session == nil { + return nil, s.closedError("graph.rag.index.records") + } + var iterator *C.CstxRagRecordIterator + if err := statusCall("graph.rag.index.records", func(errBuf *C.CstxBuffer) C.CstxStatusCode { + return C.cstx_rag_index_session_records(s.session, &iterator, errBuf) + }); err != nil { + return nil, err + } + return newNativeRagRecordIterator(iterator), nil +} + +func (s *nativeRagIndexSession) close() { + if s.session != nil { + C.cstx_rag_index_session_close(s.session) + C.cstx_rag_index_session_free(s.session) + s.session = nil + runtime.SetFinalizer(s, nil) + } +} + +// --- record iterator ----------------------------------------------------- + +type nativeRagRecordIterator struct{ iterator *C.CstxRagRecordIterator } + +func newNativeRagRecordIterator(iterator *C.CstxRagRecordIterator) *nativeRagRecordIterator { + result := &nativeRagRecordIterator{iterator: iterator} + runtime.SetFinalizer(result, (*nativeRagRecordIterator).close) + return result +} + +func (i *nativeRagRecordIterator) next(_ context.Context) (*cstxproto.RagRecord, bool, error) { + if i.iterator == nil { + return nil, false, &Error{Code: CodeInvalidArgument, Operation: "graph.rag.index.records.next", Message: "iterator is closed"} + } + var out, errBuf C.CstxBuffer + var hasValue C.uint8_t + if err := statusError( + C.cstx_rag_record_iterator_next(i.iterator, &out, &hasValue, &errBuf), + "graph.rag.index.records.next", + &errBuf, + ); err != nil { + C.cstx_buffer_free(&out) + return nil, false, err + } + data := takeBuffer(&out) + if hasValue == 0 { + return nil, false, nil + } + var record cstxproto.RagRecord + if err := proto.Unmarshal(data, &record); err != nil { + return nil, false, fmt.Errorf("cstx: decode rag record protobuf: %w", err) + } + return &record, true, nil +} + +func (i *nativeRagRecordIterator) close() { + if i.iterator != nil { + C.cstx_rag_record_iterator_close(i.iterator) + C.cstx_rag_record_iterator_free(i.iterator) + i.iterator = nil + runtime.SetFinalizer(i, nil) + } +} + +// --- retrieval ----------------------------------------------------------- + +type nativeRagRetrieval struct{ retrieval *C.CstxRagRetrieval } + +func newNativeRagRetrieval(retrieval *C.CstxRagRetrieval) *nativeRagRetrieval { + result := &nativeRagRetrieval{retrieval: retrieval} + runtime.SetFinalizer(result, (*nativeRagRetrieval).close) + return result +} + +func (r *nativeRagRetrieval) requests(_ context.Context) (*cstxproto.RecallPlan, error) { + if r.retrieval == nil { + return nil, &Error{Code: CodeNotInitialized, Operation: "graph.rag.retrieve.requests", Message: "retrieval is closed"} + } + data, err := bufferResult("graph.rag.retrieve.requests", func(out, errBuf *C.CstxBuffer) C.CstxStatusCode { + return C.cstx_rag_retrieval_requests(r.retrieval, out, errBuf) + }) + if err != nil { + return nil, err + } + var plan cstxproto.RecallPlan + if err := proto.Unmarshal(data, &plan); err != nil { + return nil, fmt.Errorf("cstx: decode recall plan protobuf: %w", err) + } + return &plan, nil +} + +func (r *nativeRagRetrieval) complete(_ context.Context, results *cstxproto.RecallResults) (*cstxproto.RagResult, error) { + if r.retrieval == nil { + return nil, &Error{Code: CodeNotInitialized, Operation: "graph.rag.complete", Message: "retrieval is closed"} + } + payload, err := proto.Marshal(results) + if err != nil { + return nil, err + } + data, err := bufferResult("graph.rag.complete", func(out, errBuf *C.CstxBuffer) C.CstxStatusCode { + rc := C.cstx_rag_retrieval_complete(r.retrieval, byteSlice(payload), out, errBuf) + runtime.KeepAlive(payload) + return rc + }) + if err != nil { + return nil, err + } + var result cstxproto.RagResult + if err := proto.Unmarshal(data, &result); err != nil { + return nil, fmt.Errorf("cstx: decode rag result protobuf: %w", err) + } + return &result, nil +} + +func (r *nativeRagRetrieval) close() { + if r.retrieval != nil { + C.cstx_rag_retrieval_close(r.retrieval) + C.cstx_rag_retrieval_free(r.retrieval) + r.retrieval = nil + runtime.SetFinalizer(r, nil) + } +} diff --git a/go/rag_test.go b/go/rag_test.go new file mode 100644 index 0000000..e30a9a4 --- /dev/null +++ b/go/rag_test.go @@ -0,0 +1,220 @@ +package cstx + +import ( + "errors" + "testing" + + "github.com/chainreactors/libcstx/go/plugins/easm" + "github.com/chainreactors/libcstx/go/proto/cstxproto" +) + +// appNode carries free text. A domain node holds only its host, which the +// projection does not turn into a searchable record, so an app is the smallest +// node that produces one. +func appNode(t *testing.T, identifier, title string) *cstxproto.Node { + t.Helper() + node, err := easm.App{AppId: identifier, Title: &title}.Node("test") + if err != nil { + t.Fatalf("build app node: %v", err) + } + id := "app:" + identifier + node.Id = &id + return node +} + +// ragRuntime builds the smallest graph that projects both a node record and a +// relationship record, which is what makes the index worth streaming. +func ragRuntime(t *testing.T) *CSTX { + t.Helper() + rt := openRuntime(t) + if _, err := rt.Graph.AddNodes(testContext, []*cstxproto.Node{ + appNode(t, "admin", "nginx admin console"), + }); err != nil { + t.Fatalf("add app node: %v", err) + } + addDomain(t, rt, "rag-target.example") + if _, err := rt.Graph.AddRelationships(testContext, []*cstxproto.Relationship{ + usesRelationship("app:admin", "domain:rag-target.example"), + }); err != nil { + t.Fatalf("add relationship: %v", err) + } + return rt +} + +func TestRagIndexStreamsProjectedRecords(t *testing.T) { + rt := ragRuntime(t) + + session, err := rt.Rag.Index(testContext, &cstxproto.RagIndexPlan{ + Commit: "revision-1", + Mode: cstxproto.RagIndexMode_RAG_INDEX_FULL, + }) + if err != nil { + t.Fatalf("index: %v", err) + } + t.Cleanup(func() { _ = session.Close() }) + + metadata, err := session.Metadata(testContext) + if err != nil { + t.Fatalf("metadata: %v", err) + } + if metadata.GetCommit() != "revision-1" { + t.Fatalf("commit = %q; want revision-1", metadata.GetCommit()) + } + if metadata.GetOperationId() == "" { + t.Fatal("projection reports no idempotent operation id") + } + if metadata.GetUpsertCount() == 0 { + t.Fatal("projection produced no records to upsert") + } + + // The iterator walks exactly the records metadata counted. + iterator, err := session.Records(testContext) + if err != nil { + t.Fatalf("records: %v", err) + } + t.Cleanup(func() { _ = iterator.Close() }) + kinds := map[cstxproto.RagRecordKind]int{} + var streamed uint64 + for { + record, ok, err := iterator.Next(testContext) + if err != nil { + t.Fatalf("next: %v", err) + } + if !ok { + break + } + if record.GetId() == "" { + t.Fatalf("record carries no id: %+v", record) + } + kinds[record.GetKind()]++ + streamed++ + } + if streamed != metadata.GetUpsertCount() { + t.Fatalf("streamed %d record(s); metadata counted %d", streamed, metadata.GetUpsertCount()) + } + if kinds[cstxproto.RagRecordKind_RAG_RECORD_NODE] == 0 { + t.Fatalf("no node records projected: %v", kinds) + } + if kinds[cstxproto.RagRecordKind_RAG_RECORD_RELATIONSHIP] == 0 { + t.Fatalf("no relationship records projected: %v", kinds) + } + + // Pending is the same projection through a caller-controlled window. + page, err := session.Pending(testContext, 0, 1) + if err != nil { + t.Fatalf("pending: %v", err) + } + if len(page.GetRecords()) != 1 { + t.Fatalf("pending returned %d record(s); want 1", len(page.GetRecords())) + } + + if _, err := session.Deletes(testContext); err != nil { + t.Fatalf("deletes: %v", err) + } +} + +func TestRagIndexSessionRejectsUseAfterClose(t *testing.T) { + rt := ragRuntime(t) + session, err := rt.Rag.Index(testContext, &cstxproto.RagIndexPlan{ + Commit: "revision-1", + Mode: cstxproto.RagIndexMode_RAG_INDEX_FULL, + }) + if err != nil { + t.Fatalf("index: %v", err) + } + if err := session.Close(); err != nil { + t.Fatalf("close: %v", err) + } + if !session.Closed() { + t.Fatal("session does not report itself closed") + } + // Repeated closes are safe, and every read after one is an error rather + // than a use of a freed native pointer. + if err := session.Close(); err != nil { + t.Fatalf("second close: %v", err) + } + if _, err := session.Deletes(testContext); err == nil { + t.Fatal("deletes succeeded on a closed session") + } + if _, err := session.Metadata(testContext); err == nil { + t.Fatal("metadata succeeded on a closed session") + } +} + +func TestRagRetrieveCompletesOnceAgainstIndexedGraph(t *testing.T) { + rt := ragRuntime(t) + if _, err := rt.Rag.Index(testContext, &cstxproto.RagIndexPlan{ + Commit: "revision-1", + Mode: cstxproto.RagIndexMode_RAG_INDEX_FULL, + }); err != nil { + t.Fatalf("index: %v", err) + } + + retrieval, err := rt.Rag.Retrieve(testContext, &cstxproto.RagQuery{ + Text: "nginx admin console", + Limit: 20, + }) + if err != nil { + t.Fatalf("retrieve: %v", err) + } + t.Cleanup(func() { _ = retrieval.Close() }) + + if _, err := retrieval.Requests(testContext); err != nil { + t.Fatalf("requests: %v", err) + } + + // A nil results means "no external recall ran", which is the built-in + // lexical case: that recall never leaves Rust. + result, err := retrieval.Complete(testContext, nil) + if err != nil { + t.Fatalf("complete: %v", err) + } + if len(result.GetNodes()) == 0 { + t.Fatalf("retrieval matched no nodes: %+v", result) + } + + // Completing consumes the retrieval. + if _, err := retrieval.Complete(testContext, nil); err == nil { + t.Fatal("retrieval completed twice") + } +} + +func TestRagRetrievalIsInvalidatedByGraphMutation(t *testing.T) { + rt := ragRuntime(t) + if _, err := rt.Rag.Index(testContext, &cstxproto.RagIndexPlan{ + Commit: "revision-1", + Mode: cstxproto.RagIndexMode_RAG_INDEX_FULL, + }); err != nil { + t.Fatalf("index: %v", err) + } + + retrieval, err := rt.Rag.Retrieve(testContext, &cstxproto.RagQuery{Text: "nginx admin console", Limit: 20}) + if err != nil { + t.Fatalf("retrieve: %v", err) + } + t.Cleanup(func() { _ = retrieval.Close() }) + + // A retrieval is bound to one graph generation, so a write between + // suspending and completing invalidates it rather than silently answering + // from a stale graph. + addDomain(t, rt, "rag-late.example") + + _, err = retrieval.Complete(testContext, nil) + if err == nil { + t.Fatal("retrieval completed against a mutated graph") + } + var cstxErr *Error + if !errors.As(err, &cstxErr) || cstxErr.Code != CodeStaleOperation { + t.Fatalf("complete error = %v; want CodeStaleOperation", err) + } +} + +func TestRagRejectsNilRequests(t *testing.T) { + rt := ragRuntime(t) + if _, err := rt.Rag.Index(testContext, nil); err == nil { + t.Fatal("index accepted a nil plan") + } + if _, err := rt.Rag.Retrieve(testContext, nil); err == nil { + t.Fatal("retrieve accepted a nil query") + } +} diff --git a/include/cstx_ffi.h b/include/cstx_ffi.h index ae52695..11fef2f 100644 --- a/include/cstx_ffi.h +++ b/include/cstx_ffi.h @@ -141,6 +141,21 @@ CstxStatusCode cstx_extension_anchor_concepts(struct CstxHandle *handle, struct CstxBuffer *output, struct CstxBuffer *error); +/** + * Parse one artifact through a native extension without mutating this graph. + * + * Input is a `ParserPayload`; output is a `Graph` batch plus the number of + * records parsed. The batch enters the same merge/link path as a parser + * implemented in any other language, so callers feed it to + * `cstx_graph_add_nodes` and then `cstx_graph_link`. The extension that owns + * the artifact must be enabled first or this reports `CSTX_NOT_FOUND`. + */ +CstxStatusCode cstx_graph_parse(struct CstxHandle *handle, + struct CstxSlice payload, + uint64_t *records, + struct CstxBuffer *output, + struct CstxBuffer *error); + /** * Add or merge a protobuf graph aggregate at the Rust-owned semantic boundary. */ diff --git a/ts/wasm/cstx_wasm_bg.wasm b/ts/wasm/cstx_wasm_bg.wasm index f8c31d0..2d35604 100644 Binary files a/ts/wasm/cstx_wasm_bg.wasm and b/ts/wasm/cstx_wasm_bg.wasm differ