mirror of
https://github.com/vxcontrol/pentagi.git
synced 2026-08-26 04:56:36 +00:00
fix(agent): stop network and malformed-arg failures from hard-failing tools
Graphiti transport errors, an empty file path, and an omitted 'action' or double-encoded 'questions' arg from the LLM all hard-failed the tool chain and burned retries instead of degrading gracefully. Also: stopTaskTimeout 5s->60s (flow kept running after a false 500), and the routine "cookie claim invalid" case now logs at Warn instead of Error.
This commit is contained in:
@@ -0,0 +1,57 @@
|
||||
package worker
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"pentagi/pkg/tools"
|
||||
)
|
||||
|
||||
// TestParseFunctionArgs_SearchInMemory_QuestionsArray verifies the CLI-args mode
|
||||
// still works for the "questions" field after it moved from []string to the more
|
||||
// lenient tools.Strings type: repeated "-questions" flags must still build a real
|
||||
// JSON array that unmarshals cleanly (the primary/common path of Strings.UnmarshalJSON).
|
||||
func TestParseFunctionArgs_SearchInMemory_QuestionsArray(t *testing.T) {
|
||||
result, err := ParseFunctionArgs(tools.SearchInMemoryToolName, []string{
|
||||
"-questions", "first query",
|
||||
"-questions", "second query",
|
||||
"-message", "test run",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("ParseFunctionArgs() unexpected error: %v", err)
|
||||
}
|
||||
|
||||
action, ok := result.(*tools.SearchInMemoryAction)
|
||||
if !ok {
|
||||
t.Fatalf("ParseFunctionArgs() returned %T, want *tools.SearchInMemoryAction", result)
|
||||
}
|
||||
|
||||
if len(action.Questions) != 2 {
|
||||
t.Fatalf("Questions length = %d, want 2 (got: %v)", len(action.Questions), []string(action.Questions))
|
||||
}
|
||||
if action.Questions[0] != "first query" || action.Questions[1] != "second query" {
|
||||
t.Errorf("Questions = %v, want [first query second query]", []string(action.Questions))
|
||||
}
|
||||
if action.Message != "test run" {
|
||||
t.Errorf("Message = %q, want %q", action.Message, "test run")
|
||||
}
|
||||
}
|
||||
|
||||
// TestParseFunctionArgs_SearchInMemory_SingleQuestion verifies a single -questions
|
||||
// flag (one array element) still round-trips correctly.
|
||||
func TestParseFunctionArgs_SearchInMemory_SingleQuestion(t *testing.T) {
|
||||
result, err := ParseFunctionArgs(tools.SearchInMemoryToolName, []string{
|
||||
"-questions", "only query",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("ParseFunctionArgs() unexpected error: %v", err)
|
||||
}
|
||||
|
||||
action, ok := result.(*tools.SearchInMemoryAction)
|
||||
if !ok {
|
||||
t.Fatalf("ParseFunctionArgs() returned %T, want *tools.SearchInMemoryAction", result)
|
||||
}
|
||||
|
||||
if len(action.Questions) != 1 || action.Questions[0] != "only query" {
|
||||
t.Fatalf("Questions = %v, want [only query]", []string(action.Questions))
|
||||
}
|
||||
}
|
||||
@@ -166,7 +166,9 @@ func InteractiveFillArgs(ctx context.Context, funcName string, taskID, subtaskID
|
||||
return structValue, nil
|
||||
}
|
||||
|
||||
// fillStructFromMap fills a structure with data from a map
|
||||
// fillStructFromMap fills a structure with data from a map. Every conversion
|
||||
// uses a safe (two-value) type assertion so a mismatch between the collected
|
||||
// value and the field's Go type returns a clear error instead of panicking.
|
||||
func fillStructFromMap(structPtr any, data map[string]any) error {
|
||||
val := reflect.ValueOf(structPtr).Elem()
|
||||
|
||||
@@ -184,32 +186,68 @@ func fillStructFromMap(structPtr any, data map[string]any) error {
|
||||
fieldName = fieldName[:comma]
|
||||
}
|
||||
|
||||
if value, ok := data[fieldName]; ok {
|
||||
fieldValue := val.Field(i)
|
||||
if fieldValue.CanSet() {
|
||||
switch fieldValue.Kind() {
|
||||
case reflect.String:
|
||||
fieldValue.SetString(value.(string))
|
||||
case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
|
||||
fieldValue.SetInt(int64(value.(int)))
|
||||
case reflect.Bool:
|
||||
fieldValue.SetBool(value.(bool))
|
||||
case reflect.Slice:
|
||||
if fieldValue.Type().Elem().Kind() == reflect.String {
|
||||
switch v := value.(type) {
|
||||
case []string:
|
||||
fieldValue.Set(reflect.ValueOf(v))
|
||||
case string:
|
||||
fieldValue.Set(reflect.ValueOf([]string{v}))
|
||||
}
|
||||
}
|
||||
case reflect.Struct:
|
||||
// For special types that may be in the tools package
|
||||
// This is a simplified version that may require refinement
|
||||
// depending on specific types
|
||||
fmt.Printf("Complex structure field detected: %s\n", fieldName)
|
||||
}
|
||||
value, ok := data[fieldName]
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
|
||||
fieldValue := val.Field(i)
|
||||
if !fieldValue.CanSet() {
|
||||
continue
|
||||
}
|
||||
|
||||
switch fieldValue.Kind() {
|
||||
case reflect.String:
|
||||
strValue, ok := value.(string)
|
||||
if !ok {
|
||||
return fmt.Errorf("unexpected type for argument '%s': expected string, got %T", fieldName, value)
|
||||
}
|
||||
fieldValue.SetString(strValue)
|
||||
|
||||
case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
|
||||
intValue, ok := value.(int)
|
||||
if !ok {
|
||||
return fmt.Errorf("unexpected type for argument '%s': expected int, got %T", fieldName, value)
|
||||
}
|
||||
fieldValue.SetInt(int64(intValue))
|
||||
|
||||
case reflect.Bool:
|
||||
boolValue, ok := value.(bool)
|
||||
if !ok {
|
||||
return fmt.Errorf("unexpected type for argument '%s': expected bool, got %T", fieldName, value)
|
||||
}
|
||||
fieldValue.SetBool(boolValue)
|
||||
|
||||
case reflect.Slice:
|
||||
if fieldValue.Type().Elem().Kind() != reflect.String {
|
||||
return fmt.Errorf("unsupported slice element type for argument '%s': %s", fieldName, fieldValue.Type().Elem().Kind())
|
||||
}
|
||||
|
||||
// Build the slice explicitly against the field's own type (e.g.
|
||||
// tools.Strings, not a plain []string) rather than relying on Go's
|
||||
// slice-assignability rules, so this keeps working correctly even
|
||||
// if the field is a named type with custom (un)marshaling.
|
||||
var items []string
|
||||
switch v := value.(type) {
|
||||
case []string:
|
||||
items = v
|
||||
case string:
|
||||
items = []string{v}
|
||||
default:
|
||||
return fmt.Errorf("unexpected type for argument '%s': expected []string or string, got %T", fieldName, value)
|
||||
}
|
||||
|
||||
converted := reflect.MakeSlice(fieldValue.Type(), len(items), len(items))
|
||||
for idx, item := range items {
|
||||
converted.Index(idx).SetString(item)
|
||||
}
|
||||
fieldValue.Set(converted)
|
||||
|
||||
case reflect.Struct:
|
||||
// For special types that may be in the tools package
|
||||
// This is a simplified version that may require refinement
|
||||
// depending on specific types
|
||||
fmt.Printf("Complex structure field detected: %s\n", fieldName)
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,73 @@
|
||||
package worker
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"pentagi/pkg/tools"
|
||||
)
|
||||
|
||||
func TestFillStructFromMap_StringsField_FromSlice(t *testing.T) {
|
||||
action := &tools.SearchInMemoryAction{}
|
||||
|
||||
err := fillStructFromMap(action, map[string]any{
|
||||
"questions": []string{"query one", "query two"},
|
||||
"message": "test",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("fillStructFromMap() unexpected error: %v", err)
|
||||
}
|
||||
if len(action.Questions) != 2 || action.Questions[0] != "query one" || action.Questions[1] != "query two" {
|
||||
t.Fatalf("Questions = %v, want [query one query two]", []string(action.Questions))
|
||||
}
|
||||
if action.Message != "test" {
|
||||
t.Fatalf("Message = %q, want %q", action.Message, "test")
|
||||
}
|
||||
}
|
||||
|
||||
func TestFillStructFromMap_StringsField_FromSingleString(t *testing.T) {
|
||||
action := &tools.SearchInMemoryAction{}
|
||||
|
||||
err := fillStructFromMap(action, map[string]any{
|
||||
"questions": "single question",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("fillStructFromMap() unexpected error: %v", err)
|
||||
}
|
||||
if len(action.Questions) != 1 || action.Questions[0] != "single question" {
|
||||
t.Fatalf("Questions = %v, want [single question]", []string(action.Questions))
|
||||
}
|
||||
}
|
||||
|
||||
func TestFillStructFromMap_TypeMismatch_ReturnsClearError_NotPanic(t *testing.T) {
|
||||
t.Run("string field given int", func(t *testing.T) {
|
||||
action := &tools.SearchInMemoryAction{}
|
||||
// The whole point of this test: a mismatched type must return an
|
||||
// error, not panic with "interface conversion: interface {} is X, not Y".
|
||||
err := fillStructFromMap(action, map[string]any{"message": 42})
|
||||
assertContainsError(t, err, "expected string")
|
||||
})
|
||||
|
||||
t.Run("int field given string", func(t *testing.T) {
|
||||
// MaxResults is a non-pointer tools.Int64 field (Kind() == Int64).
|
||||
action := &tools.SploitusAction{}
|
||||
err := fillStructFromMap(action, map[string]any{"max_results": "not-a-number"})
|
||||
assertContainsError(t, err, "expected int")
|
||||
})
|
||||
|
||||
t.Run("slice field given bool", func(t *testing.T) {
|
||||
action := &tools.SearchInMemoryAction{}
|
||||
err := fillStructFromMap(action, map[string]any{"questions": true})
|
||||
assertContainsError(t, err, "expected []string or string")
|
||||
})
|
||||
}
|
||||
|
||||
func assertContainsError(t *testing.T, err error, want string) {
|
||||
t.Helper()
|
||||
if err == nil {
|
||||
t.Fatal("expected an error for mismatched argument type, got nil")
|
||||
}
|
||||
if !strings.Contains(err.Error(), want) {
|
||||
t.Errorf("error = %q, want it to contain %q", err.Error(), want)
|
||||
}
|
||||
}
|
||||
@@ -30,7 +30,7 @@ import (
|
||||
"github.com/sirupsen/logrus"
|
||||
)
|
||||
|
||||
const stopTaskTimeout = 5 * time.Second
|
||||
const stopTaskTimeout = 60 * time.Second
|
||||
|
||||
type FlowWorker interface {
|
||||
GetFlowID() int64
|
||||
|
||||
@@ -9,6 +9,7 @@ simple:
|
||||
input: 0.30
|
||||
output: 1.20
|
||||
cache_read: 0.06
|
||||
cache_write: 0.375
|
||||
|
||||
simple_json:
|
||||
model: MiniMax-M2.7
|
||||
@@ -22,6 +23,7 @@ simple_json:
|
||||
input: 0.30
|
||||
output: 1.20
|
||||
cache_read: 0.06
|
||||
cache_write: 0.375
|
||||
|
||||
primary_agent:
|
||||
model: MiniMax-M3
|
||||
|
||||
@@ -14,6 +14,7 @@ import (
|
||||
"github.com/gin-contrib/sessions"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/jinzhu/gorm"
|
||||
"github.com/sirupsen/logrus"
|
||||
)
|
||||
|
||||
type authResult int
|
||||
@@ -73,12 +74,24 @@ func (p *AuthMiddleware) tryAuth(
|
||||
}
|
||||
|
||||
if withFail && result != authResultOk {
|
||||
response.Error(c, response.ErrAuthRequired, authErr)
|
||||
if errors.Is(authErr, errCookieClaimInvalid) {
|
||||
// An expired/absent session cookie hitting a protected endpoint is
|
||||
// routine (e.g. a stale browser tab polling after logout/expiry),
|
||||
// not an application error - log it quietly instead of at Error.
|
||||
response.ErrorWithLevel(c, response.ErrAuthRequired, authErr, logrus.WarnLevel)
|
||||
} else {
|
||||
response.Error(c, response.ErrAuthRequired, authErr)
|
||||
}
|
||||
return
|
||||
}
|
||||
c.Next()
|
||||
}
|
||||
|
||||
// errCookieClaimInvalid is returned by tryUserCookieAuthentication when the
|
||||
// session cookie is present but missing one or more required claims (expired
|
||||
// or otherwise invalid session) - a routine, expected condition.
|
||||
var errCookieClaimInvalid = errors.New("cookie claim invalid")
|
||||
|
||||
func (p *AuthMiddleware) tryUserCookieAuthentication(c *gin.Context) (authResult, error) {
|
||||
sessionObject, exists := c.Get(sessions.DefaultKey)
|
||||
if !exists {
|
||||
@@ -101,7 +114,7 @@ func (p *AuthMiddleware) tryUserCookieAuthentication(c *gin.Context) (authResult
|
||||
|
||||
for _, attr := range []any{uid, rid, prm, exp, gtm, uname, uhash, tid} {
|
||||
if attr == nil {
|
||||
return authResultFail, errors.New("cookie claim invalid")
|
||||
return authResultFail, errCookieClaimInvalid
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -37,6 +37,15 @@ func (h *HttpError) Error() string {
|
||||
}
|
||||
|
||||
func Error(c *gin.Context, err *HttpError, original error) {
|
||||
ErrorWithLevel(c, err, original, logrus.ErrorLevel)
|
||||
}
|
||||
|
||||
// ErrorWithLevel behaves exactly like Error but logs the "api error" entry at
|
||||
// the given level instead of always at Error. Use this for routes where a
|
||||
// non-2xx response is routine/expected (e.g. an unauthenticated request
|
||||
// hitting a protected endpoint) and logging it as an application error would
|
||||
// just add noise without indicating an actual problem.
|
||||
func ErrorWithLevel(c *gin.Context, err *HttpError, original error, level logrus.Level) {
|
||||
body := gin.H{
|
||||
"status": "error",
|
||||
"code": err.Code(),
|
||||
@@ -51,7 +60,7 @@ func Error(c *gin.Context, err *HttpError, original error) {
|
||||
"code": err.HttpCode(),
|
||||
"message": err.Msg(),
|
||||
}
|
||||
logger.FromContext(c).WithFields(fields).WithError(original).Error("api error")
|
||||
logger.FromContext(c).WithFields(fields).WithError(original).Log(level, "api error")
|
||||
|
||||
c.AbortWithStatusJSON(err.HttpCode(), body)
|
||||
}
|
||||
|
||||
@@ -10,6 +10,8 @@ import (
|
||||
"pentagi/pkg/version"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/sirupsen/logrus"
|
||||
"github.com/sirupsen/logrus/hooks/test"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
@@ -294,6 +296,42 @@ func TestErrorResponse_NilOriginalError(t *testing.T) {
|
||||
assert.Equal(t, "NotPermitted", body["code"])
|
||||
}
|
||||
|
||||
func TestErrorWithLevel_LogsAtGivenLevel(t *testing.T) {
|
||||
hook := test.NewGlobal()
|
||||
defer hook.Reset()
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
c.Request = httptest.NewRequest(http.MethodGet, "/test", nil)
|
||||
|
||||
ErrorWithLevel(c, ErrAuthRequired, errors.New("cookie claim invalid"), logrus.WarnLevel)
|
||||
|
||||
require.NotEmpty(t, hook.Entries)
|
||||
entry := hook.LastEntry()
|
||||
assert.Equal(t, logrus.WarnLevel, entry.Level)
|
||||
assert.Equal(t, "api error", entry.Message)
|
||||
|
||||
// The HTTP response itself must be identical to what Error() would produce -
|
||||
// only the log level differs.
|
||||
assert.Equal(t, http.StatusForbidden, w.Code)
|
||||
}
|
||||
|
||||
func TestError_StillLogsAtErrorLevel(t *testing.T) {
|
||||
hook := test.NewGlobal()
|
||||
defer hook.Reset()
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
c.Request = httptest.NewRequest(http.MethodGet, "/test", nil)
|
||||
|
||||
Error(c, ErrInternal, errors.New("db connection failed"))
|
||||
|
||||
require.NotEmpty(t, hook.Entries)
|
||||
entry := hook.LastEntry()
|
||||
assert.Equal(t, logrus.ErrorLevel, entry.Level)
|
||||
assert.Equal(t, "api error", entry.Message)
|
||||
}
|
||||
|
||||
func TestHttpError_MultipleInstancesIndependent(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
|
||||
+68
-13
@@ -1,6 +1,7 @@
|
||||
package tools
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"strconv"
|
||||
"strings"
|
||||
@@ -163,16 +164,16 @@ type MemoristResult struct {
|
||||
}
|
||||
|
||||
type SearchInMemoryAction struct {
|
||||
Questions []string `json:"questions" jsonschema:"required,minItems=1,maxItems=5" jsonschema_description:"Technical-channel payload — 1 to 5 detailed, context-rich semantic queries against the team's long-term vector store. ALWAYS written in English regardless of the engagement language: the store is indexed in English and shared across all engagements, so non-English queries will fail to retrieve relevant stored knowledge. Each query should provide context, intent, and specific details with descriptive phrases, synonyms, and related terms; multiple queries explore different semantic angles. Note: If TaskID or SubtaskID are provided, they will be used as strict filters in the search."`
|
||||
TaskID *Int64 `json:"task_id,omitempty" jsonschema:"title=Task ID" jsonschema_description:"Optional. The Task ID to use as a strict filter, retrieving information specifically related to this task. Used to enhance relevance by narrowing down the search scope. Type: integer."`
|
||||
SubtaskID *Int64 `json:"subtask_id,omitempty" jsonschema:"title=Subtask ID" jsonschema_description:"Optional. The Subtask ID to use as a strict filter, retrieving information specifically related to this subtask. Helps in refining search results for increased relevancy. Type: integer."`
|
||||
Message string `json:"message" jsonschema:"required,title=Search-in-memory message" jsonschema_description:"Engagement-log entry — a 1-2 short sentence running commentary summarizing the queries or the information retrieval process. Written in the engagement language declared by your system prompt."`
|
||||
Questions Strings `json:"questions" jsonschema:"required,type=array,minItems=1,maxItems=5" jsonschema_description:"Technical-channel payload — 1 to 5 detailed, context-rich semantic queries against the team's long-term vector store. Must be a real JSON array of strings, e.g. [\"query 1\",\"query 2\"] - NOT a JSON-encoded string containing an array. ALWAYS written in English regardless of the engagement language: the store is indexed in English and shared across all engagements, so non-English queries will fail to retrieve relevant stored knowledge. Each query should provide context, intent, and specific details with descriptive phrases, synonyms, and related terms; multiple queries explore different semantic angles. Note: If TaskID or SubtaskID are provided, they will be used as strict filters in the search."`
|
||||
TaskID *Int64 `json:"task_id,omitempty" jsonschema:"title=Task ID" jsonschema_description:"Optional. The Task ID to use as a strict filter, retrieving information specifically related to this task. Used to enhance relevance by narrowing down the search scope. Type: integer."`
|
||||
SubtaskID *Int64 `json:"subtask_id,omitempty" jsonschema:"title=Subtask ID" jsonschema_description:"Optional. The Subtask ID to use as a strict filter, retrieving information specifically related to this subtask. Helps in refining search results for increased relevancy. Type: integer."`
|
||||
Message string `json:"message" jsonschema:"required,title=Search-in-memory message" jsonschema_description:"Engagement-log entry — a 1-2 short sentence running commentary summarizing the queries or the information retrieval process. Written in the engagement language declared by your system prompt."`
|
||||
}
|
||||
|
||||
type SearchGuideAction struct {
|
||||
Questions []string `json:"questions" jsonschema:"required,minItems=1,maxItems=5" jsonschema_description:"Technical-channel payload — 1 to 5 detailed, context-rich semantic queries for the team's guide vector store. ALWAYS written in English regardless of the engagement language: the store is indexed in English and shared across all engagements, so non-English queries will fail to retrieve relevant guides. Each query should include scenario context, objectives, and specific intent. Note: The 'Type' field acts as a strict filter."`
|
||||
Type string `json:"type" jsonschema:"required,enum=install,enum=configure,enum=use,enum=pentest,enum=development,enum=other" jsonschema_description:"The specific type of guide you need. This required field acts as a strict filter to enhance the relevance of search results by narrowing down the scope to the specified guide type."`
|
||||
Message string `json:"message" jsonschema:"required,title=Guide search message" jsonschema_description:"Engagement-log entry — a 1-2 short sentence running commentary summarizing the queries and the type of guide needed. Written in the engagement language declared by your system prompt."`
|
||||
Questions Strings `json:"questions" jsonschema:"required,type=array,minItems=1,maxItems=5" jsonschema_description:"Technical-channel payload — 1 to 5 detailed, context-rich semantic queries for the team's guide vector store. Must be a real JSON array of strings, e.g. [\"query 1\",\"query 2\"] - NOT a JSON-encoded string containing an array. ALWAYS written in English regardless of the engagement language: the store is indexed in English and shared across all engagements, so non-English queries will fail to retrieve relevant guides. Each query should include scenario context, objectives, and specific intent. Note: The 'Type' field acts as a strict filter."`
|
||||
Type string `json:"type" jsonschema:"required,enum=install,enum=configure,enum=use,enum=pentest,enum=development,enum=other" jsonschema_description:"The specific type of guide you need. This required field acts as a strict filter to enhance the relevance of search results by narrowing down the scope to the specified guide type."`
|
||||
Message string `json:"message" jsonschema:"required,title=Guide search message" jsonschema_description:"Engagement-log entry — a 1-2 short sentence running commentary summarizing the queries and the type of guide needed. Written in the engagement language declared by your system prompt."`
|
||||
}
|
||||
|
||||
type StoreGuideAction struct {
|
||||
@@ -183,9 +184,9 @@ type StoreGuideAction struct {
|
||||
}
|
||||
|
||||
type SearchAnswerAction struct {
|
||||
Questions []string `json:"questions" jsonschema:"required,minItems=1,maxItems=5" jsonschema_description:"Technical-channel payload — 1 to 5 detailed, context-rich semantic queries for the team's answer vector store. ALWAYS written in English regardless of the engagement language: the store is indexed in English and shared across all engagements, so non-English queries will fail to retrieve relevant answers. Each query should include the context, what you want to find, what you intend to do with the information, and why you need it. Note: The 'Type' field acts as a strict filter."`
|
||||
Type string `json:"type" jsonschema:"required,enum=guide,enum=vulnerability,enum=code,enum=tool,enum=other" jsonschema_description:"The specific type of information or answer you are seeking. This required field acts as a strict filter to enhance the relevance of search results by narrowing down the scope to the specified type."`
|
||||
Message string `json:"message" jsonschema:"required,title=Answer search message" jsonschema_description:"Engagement-log entry — a 1-2 short sentence running commentary summarizing the queries and the type of answer needed. Written in the engagement language declared by your system prompt."`
|
||||
Questions Strings `json:"questions" jsonschema:"required,type=array,minItems=1,maxItems=5" jsonschema_description:"Technical-channel payload — 1 to 5 detailed, context-rich semantic queries for the team's answer vector store. Must be a real JSON array of strings, e.g. [\"query 1\",\"query 2\"] - NOT a JSON-encoded string containing an array. ALWAYS written in English regardless of the engagement language: the store is indexed in English and shared across all engagements, so non-English queries will fail to retrieve relevant answers. Each query should include the context, what you want to find, what you intend to do with the information, and why you need it. Note: The 'Type' field acts as a strict filter."`
|
||||
Type string `json:"type" jsonschema:"required,enum=guide,enum=vulnerability,enum=code,enum=tool,enum=other" jsonschema_description:"The specific type of information or answer you are seeking. This required field acts as a strict filter to enhance the relevance of search results by narrowing down the scope to the specified type."`
|
||||
Message string `json:"message" jsonschema:"required,title=Answer search message" jsonschema_description:"Engagement-log entry — a 1-2 short sentence running commentary summarizing the queries and the type of answer needed. Written in the engagement language declared by your system prompt."`
|
||||
}
|
||||
|
||||
type StoreAnswerAction struct {
|
||||
@@ -196,9 +197,9 @@ type StoreAnswerAction struct {
|
||||
}
|
||||
|
||||
type SearchCodeAction struct {
|
||||
Questions []string `json:"questions" jsonschema:"required,minItems=1,maxItems=5" jsonschema_description:"Technical-channel payload — 1 to 5 detailed, context-rich semantic queries for the team's code vector store. ALWAYS written in English regardless of the engagement language: the store is indexed in English and shared across all engagements, so non-English queries will fail to retrieve relevant code samples. Each query should include the context, what you intend to achieve with the code, and the functionality or content that should be included."`
|
||||
Lang string `json:"lang" jsonschema:"required" jsonschema_description:"The programming language of the code samples you need. Use the standard markdown code block language name (e.g., 'python', 'bash', 'golang'). This required field narrows down the search to code samples in the desired language."`
|
||||
Message string `json:"message" jsonschema:"required,title=Code search message" jsonschema_description:"Engagement-log entry — a 1-2 short sentence running commentary summarizing the queries and the programming language of the code samples. Written in the engagement language declared by your system prompt."`
|
||||
Questions Strings `json:"questions" jsonschema:"required,type=array,minItems=1,maxItems=5" jsonschema_description:"Technical-channel payload — 1 to 5 detailed, context-rich semantic queries for the team's code vector store. Must be a real JSON array of strings, e.g. [\"query 1\",\"query 2\"] - NOT a JSON-encoded string containing an array. ALWAYS written in English regardless of the engagement language: the store is indexed in English and shared across all engagements, so non-English queries will fail to retrieve relevant code samples. Each query should include the context, what you intend to achieve with the code, and the functionality or content that should be included."`
|
||||
Lang string `json:"lang" jsonschema:"required" jsonschema_description:"The programming language of the code samples you need. Use the standard markdown code block language name (e.g., 'python', 'bash', 'golang'). This required field narrows down the search to code samples in the desired language."`
|
||||
Message string `json:"message" jsonschema:"required,title=Code search message" jsonschema_description:"Engagement-log entry — a 1-2 short sentence running commentary summarizing the queries and the programming language of the code samples. Written in the engagement language declared by your system prompt."`
|
||||
}
|
||||
|
||||
type StoreCodeAction struct {
|
||||
@@ -400,3 +401,57 @@ func (i *Int64) String() string {
|
||||
}
|
||||
return strconv.FormatInt(int64(*i), 10)
|
||||
}
|
||||
|
||||
// Strings is a lenient []string for LLM-generated tool-call arguments (e.g.
|
||||
// the "questions" field of vector-store search tools). Models occasionally
|
||||
// double-encode the array as a JSON string containing the array literal
|
||||
// (e.g. "[\"a\", \"b\"]" instead of a real ["a","b"]), which fails a plain
|
||||
// []string unmarshal with "cannot unmarshal string into ... []string" and
|
||||
// forces an unnecessary tool-call-fixer round-trip. UnmarshalJSON recovers
|
||||
// from that case instead of failing outright.
|
||||
type Strings []string
|
||||
|
||||
func (s *Strings) UnmarshalJSON(data []byte) error {
|
||||
// A bare JSON "null" unmarshals into a nil slice with no error by default,
|
||||
// which would silently mask a required field being omitted - treat it as
|
||||
// invalid instead, consistent with Bool/Int64 above.
|
||||
if trimmed := strings.TrimSpace(string(data)); trimmed == "null" {
|
||||
return fmt.Errorf("invalid strings value: expected a JSON array of strings, got: null")
|
||||
}
|
||||
|
||||
// Primary path: a real JSON array of strings - the common, schema-conforming case.
|
||||
var arr []string
|
||||
if err := json.Unmarshal(data, &arr); err == nil {
|
||||
*s = arr
|
||||
return nil
|
||||
}
|
||||
|
||||
// Recovery path: the array was sent as a JSON string. Try to decode its
|
||||
// content as a JSON array of strings first (the double-encoding case seen
|
||||
// in production), then fall back to treating the whole string as a single
|
||||
// element so one plain question doesn't fail either.
|
||||
var encoded string
|
||||
if err := json.Unmarshal(data, &encoded); err == nil {
|
||||
var nested []string
|
||||
if err := json.Unmarshal([]byte(encoded), &nested); err == nil {
|
||||
*s = nested
|
||||
return nil
|
||||
}
|
||||
if trimmed := strings.TrimSpace(encoded); trimmed != "" {
|
||||
*s = []string{trimmed}
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
return fmt.Errorf(
|
||||
"invalid strings value: expected a JSON array of strings, e.g. [\"question 1\",\"question 2\"], got: %s",
|
||||
strings.TrimSpace(string(data)),
|
||||
)
|
||||
}
|
||||
|
||||
func (s *Strings) MarshalJSON() ([]byte, error) {
|
||||
if s == nil || *s == nil {
|
||||
return []byte("[]"), nil
|
||||
}
|
||||
return json.Marshal([]string(*s))
|
||||
}
|
||||
|
||||
@@ -592,3 +592,155 @@ func TestSearchCodeAction_QuestionsUnmarshal(t *testing.T) {
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestStringsUnmarshalJSON(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
input string
|
||||
want []string
|
||||
wantErr bool
|
||||
}{
|
||||
{
|
||||
name: "real array of strings",
|
||||
input: `["query1", "query2"]`,
|
||||
want: []string{"query1", "query2"},
|
||||
},
|
||||
{
|
||||
name: "real array single element",
|
||||
input: `["only one"]`,
|
||||
want: []string{"only one"},
|
||||
},
|
||||
{
|
||||
name: "empty array",
|
||||
input: `[]`,
|
||||
want: []string{},
|
||||
},
|
||||
{
|
||||
name: "double-encoded array (production bug case)",
|
||||
input: `"[\"SCCM MECM CM_P01\", \"ClientKeyData RSA private keys\"]"`,
|
||||
want: []string{"SCCM MECM CM_P01", "ClientKeyData RSA private keys"},
|
||||
},
|
||||
{
|
||||
name: "double-encoded single-element array",
|
||||
input: `"[\"only one\"]"`,
|
||||
want: []string{"only one"},
|
||||
},
|
||||
{
|
||||
name: "bare string falls back to single-element slice",
|
||||
input: `"just a plain question, not JSON at all"`,
|
||||
want: []string{"just a plain question, not JSON at all"},
|
||||
},
|
||||
{
|
||||
name: "empty bare string is an error",
|
||||
input: `""`,
|
||||
wantErr: true,
|
||||
},
|
||||
{
|
||||
name: "null literal is an error",
|
||||
input: `null`,
|
||||
wantErr: true,
|
||||
},
|
||||
{
|
||||
name: "array of non-strings is an error",
|
||||
input: `[1, 2, 3]`,
|
||||
wantErr: true,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
var s Strings
|
||||
err := s.UnmarshalJSON([]byte(tt.input))
|
||||
if (err != nil) != tt.wantErr {
|
||||
t.Errorf("UnmarshalJSON(%s) error = %v, wantErr %v", tt.input, err, tt.wantErr)
|
||||
return
|
||||
}
|
||||
if tt.wantErr {
|
||||
return
|
||||
}
|
||||
if len(s) != len(tt.want) {
|
||||
t.Fatalf("UnmarshalJSON(%s) = %v, want %v", tt.input, []string(s), tt.want)
|
||||
}
|
||||
for i := range tt.want {
|
||||
if s[i] != tt.want[i] {
|
||||
t.Errorf("UnmarshalJSON(%s)[%d] = %q, want %q", tt.input, i, s[i], tt.want[i])
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestStringsMarshalJSON(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
s Strings
|
||||
want string
|
||||
}{
|
||||
{name: "multiple elements", s: Strings{"a", "b"}, want: `["a","b"]`},
|
||||
{name: "single element", s: Strings{"only"}, want: `["only"]`},
|
||||
{name: "empty slice", s: Strings{}, want: "[]"},
|
||||
{name: "nil slice", s: nil, want: "[]"},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
got, err := tt.s.MarshalJSON()
|
||||
if err != nil {
|
||||
t.Fatalf("MarshalJSON() unexpected error: %v", err)
|
||||
}
|
||||
if string(got) != tt.want {
|
||||
t.Errorf("MarshalJSON() = %s, want %s", got, tt.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestStringsJSONRoundTrip(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
type container struct {
|
||||
Questions Strings `json:"questions"`
|
||||
}
|
||||
|
||||
data, err := json.Marshal(container{Questions: Strings{"q1", "q2"}})
|
||||
if err != nil {
|
||||
t.Fatalf("Marshal() error = %v", err)
|
||||
}
|
||||
|
||||
var c2 container
|
||||
if err := json.Unmarshal(data, &c2); err != nil {
|
||||
t.Fatalf("round-trip Unmarshal() error = %v", err)
|
||||
}
|
||||
if len(c2.Questions) != 2 || c2.Questions[0] != "q1" || c2.Questions[1] != "q2" {
|
||||
t.Errorf("round-trip Questions = %v, want [q1 q2]", []string(c2.Questions))
|
||||
}
|
||||
}
|
||||
|
||||
// TestSearchInMemoryAction_QuestionsDoubleEncoded reproduces the exact production
|
||||
// failure from the error log: the LLM sent "questions" as a JSON string containing
|
||||
// an escaped array literal instead of a real JSON array, which used to fail with
|
||||
// "json: cannot unmarshal string into ... []string". It must now recover gracefully.
|
||||
func TestSearchInMemoryAction_QuestionsDoubleEncoded(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
raw := `{"max_results": 10, "message": "m", "questions": "[\"SCCM MECM CM_P01 database SC_NAA Network Access Account credentials extraction\", \"SCCM ClientKeyData RSA private keys mTLS bypass MECM\"]"}`
|
||||
|
||||
var action SearchInMemoryAction
|
||||
if err := json.Unmarshal([]byte(raw), &action); err != nil {
|
||||
t.Fatalf("Unmarshal() unexpected error: %v", err)
|
||||
}
|
||||
if len(action.Questions) != 2 {
|
||||
t.Fatalf("Questions length = %d, want 2 (got: %v)", len(action.Questions), []string(action.Questions))
|
||||
}
|
||||
if action.Questions[0] != "SCCM MECM CM_P01 database SC_NAA Network Access Account credentials extraction" {
|
||||
t.Errorf("Questions[0] = %q, unexpected value", action.Questions[0])
|
||||
}
|
||||
}
|
||||
|
||||
@@ -147,6 +147,15 @@ func (b *browser) Handle(ctx context.Context, name string, args json.RawMessage)
|
||||
return "", fmt.Errorf("failed to unmarshal browser action: %w", err)
|
||||
}
|
||||
|
||||
if action.Action == "" {
|
||||
// The LLM occasionally omits the required 'action' field even though the
|
||||
// tool schema marks it required. 'markdown' is the safest default: it is
|
||||
// the most commonly used and most versatile content format, so infer it
|
||||
// instead of failing the call outright and burning a tool-call-fixer
|
||||
// round-trip on something that doesn't need one.
|
||||
action.Action = Markdown
|
||||
}
|
||||
|
||||
logger = logger.WithFields(logrus.Fields{
|
||||
"action": action.Action,
|
||||
"url": action.Url,
|
||||
@@ -163,8 +172,8 @@ func (b *browser) Handle(ctx context.Context, name string, args json.RawMessage)
|
||||
result, screen, err := b.Links(ctx, action.Url)
|
||||
return b.wrapCommandResult(ctx, name, result, action.Url, screen, err)
|
||||
default:
|
||||
logger.Error("unknown file action")
|
||||
return "", fmt.Errorf("unknown file action: %s", action.Action)
|
||||
logger.Error("unknown browser action")
|
||||
return "", fmt.Errorf("unknown browser action: %s", action.Action)
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -547,12 +547,34 @@ func TestBrowserHandle_ValidationErrors(t *testing.T) {
|
||||
|
||||
t.Run("unknown action", func(t *testing.T) {
|
||||
_, err := b.Handle(t.Context(), "browser", json.RawMessage(`{"url":"https://example.com","action":"unknown","message":"m"}`))
|
||||
if err == nil || !strings.Contains(err.Error(), "unknown file action") {
|
||||
if err == nil || !strings.Contains(err.Error(), "unknown browser action") {
|
||||
t.Fatalf("expected unknown action error, got: %v", err)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestBrowserHandle_MissingAction_DefaultsToMarkdown(t *testing.T) {
|
||||
ts := newTestScraper(t, "ok")
|
||||
defer ts.Close()
|
||||
|
||||
b := &browser{
|
||||
flowID: 1,
|
||||
dataDir: t.TempDir(),
|
||||
scPubURL: ts.URL,
|
||||
scp: &screenshotProviderMock{},
|
||||
}
|
||||
|
||||
// Mirrors production tool calls observed in the logs where the LLM omits
|
||||
// the required 'action' field entirely.
|
||||
result, err := b.Handle(t.Context(), "browser", json.RawMessage(`{"url":"https://example.com/page","message":"m"}`))
|
||||
if err != nil {
|
||||
t.Fatalf("expected inferred markdown action to succeed, got error: %v", err)
|
||||
}
|
||||
if result == "" {
|
||||
t.Fatal("Handle() returned empty result for inferred markdown action")
|
||||
}
|
||||
}
|
||||
|
||||
func TestBrowserHandle_MarkdownSuccess_StoresScreenshot(t *testing.T) {
|
||||
ts := newTestScraper(t, "ok")
|
||||
defer ts.Close()
|
||||
|
||||
@@ -3,7 +3,9 @@ package tools
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/url"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
@@ -178,6 +180,26 @@ func (t *graphitiSearchTool) Handle(ctx context.Context, name string, args json.
|
||||
}
|
||||
|
||||
if err != nil {
|
||||
// Transport-level failures (connection refused, DNS, timeout, TLS handshake
|
||||
// timeout, context deadline exceeded) surface from the underlying http.Client
|
||||
// as *url.Error. This is not a malformed-arguments problem, so it must not be
|
||||
// routed through the tool-call arg-fixer (which cannot fix a network outage
|
||||
// and would burn 3 retries doing so) — degrade gracefully instead, matching
|
||||
// the terminal/browser tools' handling of the same class of failure.
|
||||
var urlErr *url.Error
|
||||
if errors.As(err, &urlErr) {
|
||||
softMsg := fmt.Sprintf(
|
||||
"Graphiti knowledge graph is temporarily unavailable (%v); continuing without historical context.",
|
||||
err,
|
||||
)
|
||||
retriever.End(
|
||||
langfuse.WithRetrieverStatus(softMsg),
|
||||
langfuse.WithRetrieverLevel(langfuse.ObservationLevelWarning),
|
||||
)
|
||||
logger.WithError(err).Warnf("graphiti search '%s' unavailable, degrading gracefully", searchArgs.SearchType)
|
||||
return softMsg, nil
|
||||
}
|
||||
|
||||
retriever.End(
|
||||
langfuse.WithRetrieverStatus(err.Error()),
|
||||
langfuse.WithRetrieverLevel(langfuse.ObservationLevelError),
|
||||
|
||||
@@ -0,0 +1,117 @@
|
||||
package tools
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"net/url"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"pentagi/pkg/graphiti"
|
||||
)
|
||||
|
||||
// stubGraphitiSearcher is a minimal GraphitiSearcher test double: every method
|
||||
// returns whatever error/response was configured for it, so tests can exercise
|
||||
// the Handle() error-classification logic without a real Graphiti/Neo4j backend.
|
||||
type stubGraphitiSearcher struct {
|
||||
enabled bool
|
||||
err error
|
||||
}
|
||||
|
||||
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
|
||||
}
|
||||
|
||||
func (s *stubGraphitiSearcher) EntityRelationshipsSearch(
|
||||
ctx context.Context, req graphiti.EntityRelationshipSearchRequest,
|
||||
) (*graphiti.EntityRelationshipSearchResponse, error) {
|
||||
return nil, s.err
|
||||
}
|
||||
|
||||
func (s *stubGraphitiSearcher) DiverseResultsSearch(
|
||||
ctx context.Context, req graphiti.DiverseSearchRequest,
|
||||
) (*graphiti.DiverseSearchResponse, error) {
|
||||
return nil, s.err
|
||||
}
|
||||
|
||||
func (s *stubGraphitiSearcher) EpisodeContextSearch(
|
||||
ctx context.Context, req graphiti.EpisodeContextSearchRequest,
|
||||
) (*graphiti.EpisodeContextSearchResponse, error) {
|
||||
return nil, s.err
|
||||
}
|
||||
|
||||
func (s *stubGraphitiSearcher) SuccessfulToolsSearch(
|
||||
ctx context.Context, req graphiti.SuccessfulToolsSearchRequest,
|
||||
) (*graphiti.SuccessfulToolsSearchResponse, error) {
|
||||
return nil, s.err
|
||||
}
|
||||
|
||||
func (s *stubGraphitiSearcher) RecentContextSearch(
|
||||
ctx context.Context, req graphiti.RecentContextSearchRequest,
|
||||
) (*graphiti.RecentContextSearchResponse, error) {
|
||||
return nil, s.err
|
||||
}
|
||||
|
||||
func (s *stubGraphitiSearcher) EntityByLabelSearch(
|
||||
ctx context.Context, req graphiti.EntityByLabelSearchRequest,
|
||||
) (*graphiti.EntityByLabelSearchResponse, error) {
|
||||
return nil, s.err
|
||||
}
|
||||
|
||||
// fakeNetError mimics the *url.Error shape produced by http.Client.Do on a
|
||||
// transport-level failure (timeout, TLS handshake timeout, connection refused).
|
||||
func fakeNetError() error {
|
||||
return fmt.Errorf(
|
||||
"recent context search failed: failed to perform request: %w",
|
||||
&url.Error{
|
||||
Op: "Post",
|
||||
URL: "http://graphiti-neo4j/search/recent-context",
|
||||
Err: context.DeadlineExceeded,
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
func TestGraphitiSearchTool_Handle_NetworkFailure_DegradesGracefully(t *testing.T) {
|
||||
tool := NewGraphitiSearchTool(1, nil, nil, &stubGraphitiSearcher{enabled: true, err: fakeNetError()})
|
||||
|
||||
args := []byte(`{"search_type":"recent_context","query":"test query","message":"m"}`)
|
||||
result, err := tool.Handle(t.Context(), GraphitiSearchToolName, args)
|
||||
|
||||
if err != nil {
|
||||
t.Fatalf("expected graceful degradation (nil error) on transport failure, got error: %v", err)
|
||||
}
|
||||
if !strings.Contains(result, "temporarily unavailable") {
|
||||
t.Fatalf("expected soft-fail message about temporary unavailability, got: %q", result)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGraphitiSearchTool_Handle_ValidationError_StaysHard(t *testing.T) {
|
||||
tool := NewGraphitiSearchTool(1, nil, nil, &stubGraphitiSearcher{enabled: true})
|
||||
|
||||
// Invalid search_type never reaches the graphiti client - it is rejected by
|
||||
// Handle() itself, so this must remain a hard failure regardless of the
|
||||
// network-error leniency added for transport failures.
|
||||
args := []byte(`{"search_type":"not_a_real_type","query":"test query","message":"m"}`)
|
||||
_, err := tool.Handle(t.Context(), GraphitiSearchToolName, args)
|
||||
|
||||
if err == nil || !strings.Contains(err.Error(), "unknown search_type") {
|
||||
t.Fatalf("expected hard 'unknown search_type' error, got: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGraphitiSearchTool_Handle_InvalidRecencyWindow_StaysHard(t *testing.T) {
|
||||
tool := NewGraphitiSearchTool(1, nil, nil, &stubGraphitiSearcher{enabled: true})
|
||||
|
||||
// Argument-validation errors (not network errors) must still be treated as
|
||||
// hard failures so the tool-call arg-fixer can actually help here.
|
||||
args := []byte(`{"search_type":"recent_context","query":"test query","message":"m","recency_window":"not-a-window"}`)
|
||||
_, err := tool.Handle(t.Context(), GraphitiSearchToolName, args)
|
||||
|
||||
if err == nil || !strings.Contains(err.Error(), "invalid recency_window") {
|
||||
t.Fatalf("expected hard 'invalid recency_window' error, got: %v", err)
|
||||
}
|
||||
}
|
||||
@@ -144,6 +144,19 @@ func (t *terminal) Handle(ctx context.Context, name string, args json.RawMessage
|
||||
return "", fmt.Errorf("failed to unmarshal file action: %w", err)
|
||||
}
|
||||
|
||||
if action.Action == "" {
|
||||
// The LLM occasionally omits the required 'action' field even though the
|
||||
// tool schema marks it required. The intent is almost always unambiguous
|
||||
// from the other fields present, so infer it instead of failing the call
|
||||
// outright and burning a tool-call-fixer round-trip on something that
|
||||
// doesn't need one.
|
||||
if action.Content != "" {
|
||||
action.Action = WriteFile
|
||||
} else {
|
||||
action.Action = ReadFile
|
||||
}
|
||||
}
|
||||
|
||||
logger = logger.WithFields(logrus.Fields{
|
||||
"action": action.Action,
|
||||
"path": action.Path,
|
||||
@@ -308,6 +321,10 @@ func (t *terminal) getExecResult(ctx context.Context, id string, timeout time.Du
|
||||
}
|
||||
|
||||
func (t *terminal) ReadFile(ctx context.Context, flowID int64, path string) (string, error) {
|
||||
if path == "" {
|
||||
return "", fmt.Errorf("path is required and cannot be empty")
|
||||
}
|
||||
|
||||
containerName := PrimaryTerminalName(flowID)
|
||||
|
||||
isRunning, err := t.dockerClient.IsContainerRunning(ctx, t.containerLID)
|
||||
@@ -390,6 +407,10 @@ func (t *terminal) ReadFile(ctx context.Context, flowID int64, path string) (str
|
||||
}
|
||||
|
||||
func (t *terminal) WriteFile(ctx context.Context, flowID int64, content string, path string) (string, error) {
|
||||
if path == "" {
|
||||
return "", fmt.Errorf("path is required and cannot be empty")
|
||||
}
|
||||
|
||||
containerName := PrimaryTerminalName(flowID)
|
||||
|
||||
isRunning, err := t.dockerClient.IsContainerRunning(ctx, t.containerLID)
|
||||
|
||||
@@ -5,11 +5,13 @@ import (
|
||||
"bufio"
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
@@ -33,7 +35,10 @@ func (m *contextTestTermLogProvider) PutMsg(_ context.Context, _ database.Termlo
|
||||
var _ TermLogProvider = (*contextTestTermLogProvider)(nil)
|
||||
|
||||
// contextAwareMockDockerClient tracks whether the context was canceled
|
||||
// when getExecResult runs, proving context.WithoutCancel works.
|
||||
// when getExecResult runs, proving context.WithoutCancel works. It also
|
||||
// records which of CopyFromContainer (read_file) / CopyToContainer
|
||||
// (write_file) was invoked, for tests asserting on the file-tool's
|
||||
// inferred/validated action.
|
||||
type contextAwareMockDockerClient struct {
|
||||
isRunning bool
|
||||
execCreateResp container.ExecCreateResponse
|
||||
@@ -43,6 +48,10 @@ type contextAwareMockDockerClient struct {
|
||||
|
||||
// Set by ContainerExecAttach to track if ctx was canceled during attach
|
||||
ctxWasCanceled bool
|
||||
|
||||
// Set by CopyFromContainer/CopyToContainer to track which file operation ran
|
||||
copyFromCalled bool
|
||||
copyToCalled bool
|
||||
}
|
||||
|
||||
func (m *contextAwareMockDockerClient) RunContainer(_ context.Context, _ string, _ database.ContainerType,
|
||||
@@ -103,10 +112,12 @@ func (m *contextAwareMockDockerClient) ContainerExecInspect(_ context.Context, _
|
||||
return m.inspectResp, nil
|
||||
}
|
||||
func (m *contextAwareMockDockerClient) CopyToContainer(_ context.Context, _ string, _ string, _ io.Reader, _ container.CopyToContainerOptions) error {
|
||||
m.copyToCalled = true
|
||||
return nil
|
||||
}
|
||||
func (m *contextAwareMockDockerClient) CopyFromContainer(_ context.Context, _ string, _ string) (io.ReadCloser, container.PathStat, error) {
|
||||
return io.NopCloser(nil), container.PathStat{}, nil
|
||||
m.copyFromCalled = true
|
||||
return io.NopCloser(bytes.NewReader(nil)), container.PathStat{}, nil
|
||||
}
|
||||
func (m *contextAwareMockDockerClient) Cleanup(_ context.Context) error { return nil }
|
||||
func (m *contextAwareMockDockerClient) GetDefaultImage() string { return "test-image" }
|
||||
@@ -197,6 +208,127 @@ func TestExecCommandNonDetachRespectsParentCancel(t *testing.T) {
|
||||
"non-detached command SHOULD see parent context cancellation")
|
||||
}
|
||||
|
||||
func TestTerminalHandle_FileAction_DefaultsToWriteFile_WhenContentPresent(t *testing.T) {
|
||||
mock := &contextAwareMockDockerClient{isRunning: true}
|
||||
term := &terminal{
|
||||
flowID: 1,
|
||||
containerID: 1,
|
||||
containerLID: "test-container",
|
||||
dockerClient: mock,
|
||||
tlp: &contextTestTermLogProvider{},
|
||||
}
|
||||
|
||||
// No "action" field at all - mirrors the malformed tool calls observed in
|
||||
// production logs (LLM omits the required 'action' when content+path make
|
||||
// the intent unambiguous).
|
||||
args := json.RawMessage(`{"path":"/work/test.py","content":"print(1)","message":"m"}`)
|
||||
_, err := term.Handle(t.Context(), FileToolName, args)
|
||||
|
||||
if err != nil {
|
||||
t.Fatalf("expected inferred write_file to succeed, got error: %v", err)
|
||||
}
|
||||
if !mock.copyToCalled || mock.copyFromCalled {
|
||||
t.Fatalf("expected CopyToContainer (write_file) to be called, copyTo=%v copyFrom=%v",
|
||||
mock.copyToCalled, mock.copyFromCalled)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTerminalHandle_FileAction_DefaultsToReadFile_WhenContentAbsent(t *testing.T) {
|
||||
mock := &contextAwareMockDockerClient{isRunning: true}
|
||||
term := &terminal{
|
||||
flowID: 1,
|
||||
containerID: 1,
|
||||
containerLID: "test-container",
|
||||
dockerClient: mock,
|
||||
tlp: &contextTestTermLogProvider{},
|
||||
}
|
||||
|
||||
args := json.RawMessage(`{"path":"/work/test.py","message":"m"}`)
|
||||
_, err := term.Handle(t.Context(), FileToolName, args)
|
||||
|
||||
if err != nil {
|
||||
t.Fatalf("expected inferred read_file to succeed, got error: %v", err)
|
||||
}
|
||||
if !mock.copyFromCalled || mock.copyToCalled {
|
||||
t.Fatalf("expected CopyFromContainer (read_file) to be called, copyFrom=%v copyTo=%v",
|
||||
mock.copyFromCalled, mock.copyToCalled)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTerminalHandle_FileAction_EmptyPath_ReadFile_ReturnsClearError(t *testing.T) {
|
||||
mock := &contextAwareMockDockerClient{isRunning: true}
|
||||
term := &terminal{
|
||||
flowID: 1,
|
||||
containerID: 1,
|
||||
containerLID: "test-container",
|
||||
dockerClient: mock,
|
||||
tlp: &contextTestTermLogProvider{},
|
||||
}
|
||||
|
||||
args := json.RawMessage(`{"action":"read_file","path":"","message":"m"}`)
|
||||
// Handle() wraps ReadFile/WriteFile errors into a soft (nil-error) response
|
||||
// via wrapCommandResult, same as every other terminal-tool failure - so the
|
||||
// error text is asserted on the returned string, not a returned Go error.
|
||||
result, err := term.Handle(t.Context(), FileToolName, args)
|
||||
|
||||
if err != nil {
|
||||
t.Fatalf("expected soft-failed (nil error) response, got error: %v", err)
|
||||
}
|
||||
if !strings.Contains(result, "path is required and cannot be empty") {
|
||||
t.Fatalf("expected result to mention the empty-path error, got: %q", result)
|
||||
}
|
||||
if mock.copyFromCalled {
|
||||
t.Fatal("expected CopyFromContainer to not be called for an empty path")
|
||||
}
|
||||
}
|
||||
|
||||
func TestTerminalHandle_FileAction_EmptyPath_WriteFile_ReturnsClearError(t *testing.T) {
|
||||
mock := &contextAwareMockDockerClient{isRunning: true}
|
||||
term := &terminal{
|
||||
flowID: 1,
|
||||
containerID: 1,
|
||||
containerLID: "test-container",
|
||||
dockerClient: mock,
|
||||
tlp: &contextTestTermLogProvider{},
|
||||
}
|
||||
|
||||
args := json.RawMessage(`{"action":"write_file","path":"","content":"data","message":"m"}`)
|
||||
result, err := term.Handle(t.Context(), FileToolName, args)
|
||||
|
||||
if err != nil {
|
||||
t.Fatalf("expected soft-failed (nil error) response, got error: %v", err)
|
||||
}
|
||||
if !strings.Contains(result, "path is required and cannot be empty") {
|
||||
t.Fatalf("expected result to mention the empty-path error, got: %q", result)
|
||||
}
|
||||
if mock.copyToCalled {
|
||||
t.Fatal("expected CopyToContainer to not be called for an empty path")
|
||||
}
|
||||
}
|
||||
|
||||
func TestTerminalHandle_FileAction_ExplicitInvalidAction_StillFails(t *testing.T) {
|
||||
mock := &contextAwareMockDockerClient{isRunning: true}
|
||||
term := &terminal{
|
||||
flowID: 1,
|
||||
containerID: 1,
|
||||
containerLID: "test-container",
|
||||
dockerClient: mock,
|
||||
tlp: &contextTestTermLogProvider{},
|
||||
}
|
||||
|
||||
// An explicit but invalid action must still be a hard failure - inference
|
||||
// only kicks in when the field is empty.
|
||||
args := json.RawMessage(`{"path":"/work/test.py","action":"delete_file","message":"m"}`)
|
||||
_, err := term.Handle(t.Context(), FileToolName, args)
|
||||
|
||||
if err == nil || !strings.Contains(err.Error(), "unknown file action") {
|
||||
t.Fatalf("expected unknown file action error, got: %v", err)
|
||||
}
|
||||
if mock.copyFromCalled || mock.copyToCalled {
|
||||
t.Fatalf("expected no docker calls for an invalid explicit action")
|
||||
}
|
||||
}
|
||||
|
||||
func TestPrimaryTerminalName(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
|
||||
Reference in New Issue
Block a user