diff --git a/backend/pkg/server/services/assistants.go b/backend/pkg/server/services/assistants.go index 2000934c..71c6672c 100644 --- a/backend/pkg/server/services/assistants.go +++ b/backend/pkg/server/services/assistants.go @@ -20,6 +20,7 @@ import ( "github.com/gin-gonic/gin" "github.com/jinzhu/gorm" + "github.com/sirupsen/logrus" ) type assistants struct { @@ -432,6 +433,10 @@ func (s *AssistantService) PatchAssistant(c *gin.Context) { fw, err := s.fc.GetFlow(c, int64(flowID)) if err != nil { + if errors.Is(err, controller.ErrFlowNotFound) { + response.ErrorWithLevel(c, response.ErrFlowsNotFound, err, logrus.WarnLevel) + return + } logger.FromContext(c).WithError(err).Errorf("error getting flow by id in flow controller") response.Error(c, response.ErrInternal, err) return @@ -565,6 +570,10 @@ func (s *AssistantService) DeleteAssistant(c *gin.Context) { fw, err := s.fc.GetFlow(c, int64(flowID)) if err != nil { + if errors.Is(err, controller.ErrFlowNotFound) { + response.ErrorWithLevel(c, response.ErrFlowsNotFound, err, logrus.WarnLevel) + return + } logger.FromContext(c).WithError(err).Errorf("error getting flow by id in flow controller") response.Error(c, response.ErrInternal, err) return diff --git a/backend/pkg/server/services/flows.go b/backend/pkg/server/services/flows.go index 03fcb5ae..374d6cf9 100644 --- a/backend/pkg/server/services/flows.go +++ b/backend/pkg/server/services/flows.go @@ -19,6 +19,7 @@ import ( "github.com/gin-gonic/gin" "github.com/jinzhu/gorm" + "github.com/sirupsen/logrus" ) type flows struct { @@ -450,6 +451,10 @@ func (s *FlowService) PatchFlow(c *gin.Context) { fw, err := s.fc.GetFlow(c, int64(flow.ID)) if err != nil { + if errors.Is(err, controller.ErrFlowNotFound) { + response.ErrorWithLevel(c, response.ErrFlowsNotFound, err, logrus.WarnLevel) + return + } logger.FromContext(c).WithError(err).Errorf("error getting flow by id in flow controller") response.Error(c, response.ErrInternal, err) return diff --git a/backend/pkg/templates/prompts/enricher.tmpl b/backend/pkg/templates/prompts/enricher.tmpl index f7633391..5e165dc6 100644 --- a/backend/pkg/templates/prompts/enricher.tmpl +++ b/backend/pkg/templates/prompts/enricher.tmpl @@ -171,7 +171,7 @@ Follow this prioritized approach to gather SUPPLEMENTARY information: Search knowledge graph for episodic memory and execution history Find what agents discovered and executed during operations -recent_context, episode_context, successful_tools, entity_relationships (only after another search type returns a node UUID; needs center_node_uuid) +recent_context, episode_context, successful_tools, entity_relationships (needs center_node_uuid copied verbatim from a 'UUID:' field in an EARLIER result of this conversation — never invent one, e.g. not a flow ID or hostname) {{- end}} diff --git a/backend/pkg/templates/prompts/memorist.tmpl b/backend/pkg/templates/prompts/memorist.tmpl index 18993ee6..d1442239 100644 --- a/backend/pkg/templates/prompts/memorist.tmpl +++ b/backend/pkg/templates/prompts/memorist.tmpl @@ -76,7 +76,7 @@ Choose the appropriate search type based on the information need: - When: Looking for working command examples, successful approaches - Example: `search_type: "successful_tools", query: "successful nmap scans revealing services", min_mentions: 2` -5. **entity_relationships** - Explore entity connections (requires center_node_uuid from a prior search result; node_labels/edge_types are optional filters from the taxonomy reference above) +5. **entity_relationships** - Explore entity connections (requires center_node_uuid copied verbatim from a 'UUID:' field in a prior search result — never invent one; node_labels/edge_types are optional filters from the taxonomy reference above) - Use: "What is connected to [entity]?" - When: Understanding relationships between discovered entities - Example: `search_type: "entity_relationships", query: "related vulnerabilities and services", center_node_uuid: "[uuid]", max_depth: 2` diff --git a/backend/pkg/templates/prompts/pentester.tmpl b/backend/pkg/templates/prompts/pentester.tmpl index 7404fc11..db441257 100644 --- a/backend/pkg/templates/prompts/pentester.tmpl +++ b/backend/pkg/templates/prompts/pentester.tmpl @@ -93,7 +93,7 @@ Choose the appropriate search type based on your need: - When: Need detailed context, understanding decision-making - Example: `search_type: "episode_context", query: "pentester agent analysis of SSH vulnerability"` -4. **entity_relationships** - Explore entity connections (requires center_node_uuid from a prior search result; node_labels/edge_types are optional filters from the taxonomy reference above) +4. **entity_relationships** - Explore entity connections (requires center_node_uuid copied verbatim from a 'UUID:' field in a prior search result — never invent one; node_labels/edge_types are optional filters from the taxonomy reference above) - Use: "What services/vulnerabilities are related to [entity]?" - When: Investigating a specific IP, service, or vulnerability - Example: `search_type: "entity_relationships", query: "services and vulnerabilities", center_node_uuid: "[uuid]", max_depth: 2` diff --git a/backend/pkg/tools/args.go b/backend/pkg/tools/args.go index 6618d679..2f90ff3a 100644 --- a/backend/pkg/tools/args.go +++ b/backend/pkg/tools/args.go @@ -160,11 +160,11 @@ type WebSearchAction struct { type GraphitiSearchAction struct { SearchType String `json:"search_type" jsonschema:"required,type=string,enum=temporal_window,enum=entity_relationships,enum=diverse_results,enum=episode_context,enum=successful_tools,enum=recent_context,enum=entity_by_label" jsonschema_description:"Type of search to perform: temporal_window (time-bounded search), entity_relationships (graph traversal from an entity), diverse_results (anti-redundancy search), episode_context (full agent reasoning and tool outputs), successful_tools (proven techniques), recent_context (latest findings), entity_by_label (type-specific entity search, REQUIRES node_labels)"` - Query string `json:"query" jsonschema:"required" jsonschema_description:"Technical-channel payload — natural language query against the team's temporal knowledge graph. ALWAYS written in English regardless of the engagement language: the graph is indexed in English and shared across all engagements; non-English queries will fail to retrieve stored episodic memory."` + Query string `json:"query" jsonschema:"required" jsonschema_description:"Technical-channel payload — natural language query against the team's temporal knowledge graph. REQUIRED for EVERY search_type, including entity_relationships/entity_by_label: it re-ranks/filters the graph traversal, it is never optional just because center_node_uuid or node_labels is set. ALWAYS written in English regardless of the engagement language: the graph is indexed in English and shared across all engagements; non-English queries will fail to retrieve stored episodic memory."` MaxResults *Int64 `json:"max_results,omitempty" jsonschema:"title=Maximum Results,type=integer" jsonschema_description:"Maximum number of results to return (default varies by search type)"` TimeStart string `json:"time_start,omitempty" jsonschema_description:"Start of time window (ISO 8601 format, required for temporal_window)"` TimeEnd string `json:"time_end,omitempty" jsonschema_description:"End of time window (ISO 8601 format, required for temporal_window)"` - CenterNodeUUID string `json:"center_node_uuid,omitempty" jsonschema_description:"UUID of entity to search from (required for entity_relationships)"` + CenterNodeUUID string `json:"center_node_uuid,omitempty" jsonschema_description:"REQUIRED for entity_relationships. Copy verbatim from the 'UUID:' field of an entity/community returned by an EARLIER graphiti_search call in this conversation (any search_type). NEVER invent one — not a flow/task ID, hostname, or title; if no prior result yielded a UUID, run recent_context/entity_by_label first"` MaxDepth *Int64 `json:"max_depth,omitempty" jsonschema:"title=Maximum Depth,type=integer" jsonschema_description:"Maximum graph traversal depth (default: 2, max: 3, for entity_relationships)"` NodeLabels []string `json:"node_labels,omitempty" jsonschema_description:"Filter to specific node types — EXACT taxonomy names, PascalCase singular (e.g., ['Host', 'Service', 'Vulnerability']). REQUIRED (non-empty) when search_type is entity_by_label; optional filter otherwise"` EdgeTypes []string `json:"edge_types,omitempty" jsonschema_description:"Filter to specific relationship types — EXACT taxonomy names, UPPER_SNAKE_CASE (e.g., ['HAS_PORT', 'HAS_VULNERABILITY'])"` @@ -387,6 +387,10 @@ type Int64 int64 func (i *Int64) UnmarshalJSON(data []byte) error { sdata := strings.Trim(strings.ToLower(string(data)), "' \"\n\r\t") + if sdata == "" { + *i = 0 + return nil + } num, err := strconv.ParseInt(sdata, 10, 64) if err != nil { return fmt.Errorf("invalid int value: %s", sdata) diff --git a/backend/pkg/tools/args_test.go b/backend/pkg/tools/args_test.go index c7a68458..391784d5 100644 --- a/backend/pkg/tools/args_test.go +++ b/backend/pkg/tools/args_test.go @@ -169,7 +169,7 @@ func TestInt64UnmarshalJSON(t *testing.T) { {name: "underflow int64", input: `"-9223372036854775809"`, wantErr: true}, {name: "invalid string", input: `"abc"`, wantErr: true}, {name: "invalid float", input: `"1.5"`, wantErr: true}, - {name: "empty string", input: `""`, wantErr: true}, + {name: "empty string treated as 0 (production near-miss: LLM sends timeout: \"\")", input: `""`, want: 0}, {name: "bool string", input: `"true"`, wantErr: true}, } @@ -422,6 +422,20 @@ func TestInt64JSONRoundTrip(t *testing.T) { } } +func TestTerminalAction_EmptyTimeout_ProductionBugReproduction(t *testing.T) { + t.Parallel() + + data := []byte(`{"input": "curl -s http://example.com/", "cwd": "/", "detach": false, "message": "test", "timeout": ""}`) + + var action TerminalAction + if err := json.Unmarshal(data, &action); err != nil { + t.Fatalf("Unmarshal() unexpected error for timeout: \"\": %v", err) + } + if action.Timeout != 0 { + t.Errorf("Timeout = %v, want 0 (empty string treated as unset/default)", action.Timeout) + } +} + func TestSearchInMemoryAction_QuestionsUnmarshal(t *testing.T) { t.Parallel() diff --git a/backend/pkg/tools/graphiti_search.go b/backend/pkg/tools/graphiti_search.go index 8dc10d9c..7e5c8b8d 100644 --- a/backend/pkg/tools/graphiti_search.go +++ b/backend/pkg/tools/graphiti_search.go @@ -58,6 +58,35 @@ const ( DefaultRecencyWindow = "24h" ) +// graphitiTimeFormats are the layouts accepted for temporal_window's +// time_start/time_end, tried in order. RFC3339 is the documented format; +// the remaining layouts tolerate a common LLM near-miss - a timestamp +// missing the timezone designator (e.g. "2026-07-24T11:53:34") - which +// would otherwise fail with a confusing parse error even though the value +// is unambiguous. Layouts without a zone parse as UTC, matching the +// graph's own UTC timestamps. +var graphitiTimeFormats = []string{ + time.RFC3339, + "2006-01-02T15:04:05", + "2006-01-02 15:04:05", +} + +// parseGraphitiTime parses a temporal_window time_start/time_end value, +// trying each of graphitiTimeFormats in order and returning the error from +// the strict RFC3339 attempt if none match (the most informative for the LLM +// tool-call fixer, since it's the documented/expected format). +func parseGraphitiTime(value string) (time.Time, error) { + rfc3339Err := error(nil) + for i, format := range graphitiTimeFormats { + if t, err := time.Parse(format, value); err == nil { + return t, nil + } else if i == 0 { + rfc3339Err = err + } + } + return time.Time{}, rfc3339Err +} + var ( allowedRecencyWindows = map[string]struct{}{ "1h": {}, @@ -316,14 +345,14 @@ func (t *graphitiSearchTool) handleTemporalWindowSearch( return "", fmt.Errorf("time_start and time_end are required for temporal_window search") } - timeStart, err := time.Parse(time.RFC3339, args.TimeStart) + timeStart, err := parseGraphitiTime(args.TimeStart) if err != nil { - return "", fmt.Errorf("invalid time_start format (use ISO 8601): %w", err) + return "", fmt.Errorf("invalid time_start format (use ISO 8601, e.g. 2026-01-02T15:04:05Z): %w", err) } - timeEnd, err := time.Parse(time.RFC3339, args.TimeEnd) + timeEnd, err := parseGraphitiTime(args.TimeEnd) if err != nil { - return "", fmt.Errorf("invalid time_end format (use ISO 8601): %w", err) + return "", fmt.Errorf("invalid time_end format (use ISO 8601, e.g. 2026-01-02T15:04:05Z): %w", err) } if timeEnd.Before(timeStart) { @@ -728,6 +757,7 @@ func FormatGraphitiDiverseResults( score = fmt.Sprintf(" (MMR score: %.3f)", resp.CommunityMMRScores[i]) } builder.WriteString(fmt.Sprintf("%d. **%s**%s\n", i+1, comm.Name, score)) + builder.WriteString(fmt.Sprintf(" - UUID: %s\n", comm.UUID)) builder.WriteString(fmt.Sprintf(" - Summary: %s\n\n", comm.Summary)) } } @@ -791,7 +821,7 @@ func FormatGraphitiEpisodeContextResults( if i < len(resp.MentionedNodeScores) { score = fmt.Sprintf(" (relevance: %.3f)", resp.MentionedNodeScores[i]) } - builder.WriteString(fmt.Sprintf("- **%s**%s: %s\n", node.Name, score, node.Summary)) + builder.WriteString(fmt.Sprintf("- **%s**%s (UUID: %s): %s\n", node.Name, score, node.UUID, node.Summary)) } } @@ -864,6 +894,7 @@ func FormatGraphitiRecentContextResults( score = fmt.Sprintf(" (score: %.3f)", resp.NodeScores[i]) } builder.WriteString(fmt.Sprintf("%d. **%s**%s\n", i+1, node.Name, score)) + builder.WriteString(fmt.Sprintf(" - UUID: %s\n", node.UUID)) builder.WriteString(fmt.Sprintf(" - Labels: %v\n", node.Labels)) builder.WriteString(fmt.Sprintf(" - Summary: %s\n\n", node.Summary)) } diff --git a/backend/pkg/tools/graphiti_search_test.go b/backend/pkg/tools/graphiti_search_test.go index 4001ac34..5ac93e3e 100644 --- a/backend/pkg/tools/graphiti_search_test.go +++ b/backend/pkg/tools/graphiti_search_test.go @@ -2,8 +2,11 @@ package tools import ( "context" + "encoding/json" "fmt" + "maps" "net/url" + "slices" "strings" "testing" @@ -23,7 +26,10 @@ func (s *stubGraphitiSearcher) IsEnabled() bool { return s.enabled } func (s *stubGraphitiSearcher) TemporalWindowSearch( ctx context.Context, req graphiti.TemporalSearchRequest, ) (*graphiti.TemporalSearchResponse, error) { - return nil, s.err + if s.err != nil { + return nil, s.err + } + return &graphiti.TemporalSearchResponse{}, nil } func (s *stubGraphitiSearcher) EntityRelationshipsSearch( @@ -251,6 +257,57 @@ func TestGraphitiSearchTool_Handle_EntityRelationships_ValidCenterNodeUUID_Succe } } +func TestParseGraphitiTime(t *testing.T) { + tests := []struct { + name string + input string + wantErr bool + }{ + {"RFC3339 with Z", "2026-07-25T11:53:34Z", false}, + {"RFC3339 with offset", "2026-07-25T11:53:34+03:00", false}, + {"missing timezone designator (production near-miss)", "2026-07-24T11:53:34", false}, + {"space-separated, no timezone", "2026-07-24 11:53:34", false}, + {"garbage", "not-a-date", true}, + {"empty", "", true}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + _, err := parseGraphitiTime(tt.input) + if tt.wantErr && err == nil { + t.Fatalf("expected an error for input %q, got nil", tt.input) + } + if !tt.wantErr && err != nil { + t.Fatalf("expected no error for input %q, got: %v", tt.input, err) + } + }) + } +} + +func TestGraphitiSearchTool_Handle_TemporalWindow_MissingTimezone_Succeeds(t *testing.T) { + tool := NewGraphitiSearchTool(1, nil, nil, &stubGraphitiSearcher{enabled: true}) + + // Reproduces the exact production near-miss: an LLM omitted the trailing + // 'Z'/offset on an otherwise well-formed timestamp. + args := []byte(`{"search_type":"temporal_window","query":"test query","time_start":"2026-07-24T11:53:34","time_end":"2026-07-25T11:53:34","message":"m"}`) + _, err := tool.Handle(t.Context(), GraphitiSearchToolName, args) + + if err != nil { + t.Fatalf("expected no error for a timestamp missing its timezone designator, got: %v", err) + } +} + +func TestGraphitiSearchTool_Handle_TemporalWindow_GarbageTime_StaysHard(t *testing.T) { + tool := NewGraphitiSearchTool(1, nil, nil, &stubGraphitiSearcher{enabled: true}) + + args := []byte(`{"search_type":"temporal_window","query":"test query","time_start":"not-a-date","time_end":"2026-07-25T11:53:34Z","message":"m"}`) + _, err := tool.Handle(t.Context(), GraphitiSearchToolName, args) + + if err == nil || !strings.Contains(err.Error(), "invalid time_start format") { + t.Fatalf("expected hard 'invalid time_start format' error, got: %v", err) + } +} + func TestGraphitiSearchTool_Handle_InvalidRecencyWindow_StaysHard(t *testing.T) { tool := NewGraphitiSearchTool(1, nil, nil, &stubGraphitiSearcher{enabled: true}) @@ -263,3 +320,155 @@ func TestGraphitiSearchTool_Handle_InvalidRecencyWindow_StaysHard(t *testing.T) t.Fatalf("expected hard 'invalid recency_window' error, got: %v", err) } } + +// These three tests are regression guards for a real production bug: the +// entity/community listings for recent_context, diverse_results, and +// episode_context silently omitted the "UUID:" field that +// FormatGraphitiTemporalResults, FormatGraphitiEntityRelationshipResults, and +// FormatGraphitiEntityByLabelResults already included. Since entity_relationships +// requires a real center_node_uuid copied from an EARLIER result, agents that +// started their research with recent_context (the documented default) never +// saw a UUID to copy and fabricated one instead (a flow ID, a hostname, a +// page title) - which is exactly what graphiti_search's new UUID-format +// validation now (correctly) rejects. Fixing only the validation without also +// restoring the missing UUID field would leave every recent_context-first +// workflow permanently unable to reach entity_relationships. + +func TestFormatGraphitiRecentContextResults_IncludesNodeUUID(t *testing.T) { + resp := &graphiti.RecentContextSearchResponse{ + Nodes: []graphiti.NodeResult{ + {UUID: "f7b95dfc-ee58-4a8b-8d85-582cf117b4df", Name: "500", Labels: []string{"Entity", "Port"}, Summary: "Host has port 500"}, + }, + NodeScores: []float64{0.5}, + } + + result := FormatGraphitiRecentContextResults(resp, "test query") + + if !strings.Contains(result, "UUID: f7b95dfc-ee58-4a8b-8d85-582cf117b4df") { + t.Fatalf("expected recent_context entity listing to include the node UUID, got:\n%s", result) + } +} + +func TestFormatGraphitiDiverseResults_IncludesCommunityUUID(t *testing.T) { + resp := &graphiti.DiverseSearchResponse{ + Communities: []graphiti.CommunityResult{ + {UUID: "a1b2c3d4-e5f6-4789-a012-3456789abcde", Name: "Trading Platform Cluster", Summary: "Cluster of related findings"}, + }, + CommunityMMRScores: []float64{0.9}, + } + + result := FormatGraphitiDiverseResults(resp, "test query") + + if !strings.Contains(result, "UUID: a1b2c3d4-e5f6-4789-a012-3456789abcde") { + t.Fatalf("expected diverse_results community listing to include the community UUID, got:\n%s", result) + } +} + +func TestFormatGraphitiEpisodeContextResults_IncludesMentionedNodeUUID(t *testing.T) { + resp := &graphiti.EpisodeContextSearchResponse{ + MentionedNodes: []graphiti.NodeResult{ + {UUID: "11111111-2222-4333-8444-555555555555", Name: "NoSQL Injection", Summary: "Confirmed vulnerability"}, + }, + MentionedNodeScores: []float64{0.8}, + } + + result := FormatGraphitiEpisodeContextResults(resp, "test query") + + if !strings.Contains(result, "UUID: 11111111-2222-4333-8444-555555555555") { + t.Fatalf("expected episode_context mentioned-entities listing to include the node UUID, got:\n%s", result) + } +} + +// TestGraphitiSearchAction_JSONSchema_IncludesAllFields serializes the exact +// JSON schema the LLM is shown for the graphiti_search tool (the same +// reflector.Reflect(&GraphitiSearchAction{}) call registry.go uses to build +// the live tool definition) and verifies every struct field survives into +// the "properties" object under its declared json name, and that "required" +// matches the struct tags exactly. +// +// This exists to answer a concrete production question: agents were sending +// hallucinated values (a flow ID, a hostname, a page title, even a stray +// markdown line) for center_node_uuid instead of a real Graphiti UUID. This +// test rules out "the LLM never actually saw the field / its guidance" as a +// cause — if it ever regresses (a field silently dropped, unexported, or +// renamed without updating the json tag), this test fails at build/test time +// instead of surfacing only as a confusing runtime tool-call error. +func TestGraphitiSearchAction_JSONSchema_IncludesAllFields(t *testing.T) { + schema := reflector.Reflect(&GraphitiSearchAction{}) + + raw, err := json.Marshal(schema) + if err != nil { + t.Fatalf("failed to marshal schema: %v", err) + } + + var doc map[string]any + if err := json.Unmarshal(raw, &doc); err != nil { + t.Fatalf("failed to unmarshal schema into a generic map: %v", err) + } + + properties, ok := doc["properties"].(map[string]any) + if !ok { + t.Fatalf("schema has no 'properties' object, got: %s", raw) + } + + // Every field of GraphitiSearchAction, by its json tag name. + wantFields := []string{ + "search_type", "query", "max_results", "time_start", "time_end", + "center_node_uuid", "max_depth", "node_labels", "edge_types", + "diversity_level", "min_mentions", "recency_window", "message", + } + for _, field := range wantFields { + if _, ok := properties[field]; !ok { + t.Errorf("expected field %q in the JSON schema 'properties' but it is missing (full schema below)\n%s", field, raw) + } + } + if len(properties) != len(wantFields) { + t.Errorf("expected exactly %d properties, got %d: %v", len(wantFields), len(properties), slices.Collect(maps.Keys(properties))) + } + + requiredRaw, ok := doc["required"].([]any) + if !ok { + t.Fatalf("schema has no 'required' array, got: %s", raw) + } + var required []string + for _, r := range requiredRaw { + s, ok := r.(string) + if !ok { + t.Fatalf("required entry is not a string: %v", r) + } + required = append(required, s) + } + wantRequired := []string{"search_type", "query", "message"} + if len(required) != len(wantRequired) { + t.Fatalf("expected required=%v, got %v", wantRequired, required) + } + for _, field := range wantRequired { + if !slices.Contains(required, field) { + t.Errorf("expected %q to be in the required list, got %v", field, required) + } + } + + // Spot-check the exact guidance an LLM would read for the two fields + // implicated in the production failure, so a future edit that weakens + // or removes this wording is caught here rather than in a live flow. + centerNodeUUID, ok := properties["center_node_uuid"].(map[string]any) + if !ok { + t.Fatalf("center_node_uuid property is not an object: %v", properties["center_node_uuid"]) + } + centerDesc, _ := centerNodeUUID["description"].(string) + if !strings.Contains(centerDesc, "NEVER invent") { + t.Errorf("expected center_node_uuid description to warn against inventing a value, got: %q", centerDesc) + } + if !strings.Contains(centerDesc, "UUID:") { + t.Errorf("expected center_node_uuid description to reference the 'UUID:' field agents must copy from, got: %q", centerDesc) + } + + nodeLabels, ok := properties["node_labels"].(map[string]any) + if !ok { + t.Fatalf("node_labels property is not an object: %v", properties["node_labels"]) + } + nodeLabelsDesc, _ := nodeLabels["description"].(string) + if !strings.Contains(nodeLabelsDesc, "PascalCase") { + t.Errorf("expected node_labels description to mention PascalCase casing, got: %q", nodeLabelsDesc) + } +}