From b234dbacd4ce688aecbc6485aeb3293475a8b9c2 Mon Sep 17 00:00:00 2001 From: riturajFi Date: Mon, 9 Feb 2026 09:04:56 +0530 Subject: [PATCH 01/20] feat(api/sdk): editing trigger + auto-trim strategies - Gate edit_strategies with editing_trigger\n- Add auto-trim strategies + trigger checks\n- Update Python/TypeScript SDK sessions + docs --- .../docs/(guides)/engineering/editing.mdx | 27 ++-- .../src/acontext/resources/async_sessions.py | 7 + .../src/acontext/resources/sessions.py | 7 + .../acontext-ts/src/resources/sessions.ts | 5 + .../go/internal/modules/handler/session.go | 55 ++++++++ .../go/internal/modules/service/session.go | 132 +++++++++++++++--- 6 files changed, 200 insertions(+), 33 deletions(-) diff --git a/docs/content/docs/(guides)/engineering/editing.mdx b/docs/content/docs/(guides)/engineering/editing.mdx index 3efc2f3e1..7cfe688d9 100644 --- a/docs/content/docs/(guides)/engineering/editing.mdx +++ b/docs/content/docs/(guides)/engineering/editing.mdx @@ -9,32 +9,25 @@ Apply edit strategies when retrieving messages to manage context window size. Th The `get_messages` response includes `this_time_tokens` - the total token count of returned messages. Use this to: - Check current context window size -- Decide when to apply edit strategies +- Apply edit strategies only when needed - Determine when to [reset the prompt cache](/engineering/cache) ```python title="Python" -result = client.sessions.get_messages(session_id="session-uuid") +result = client.sessions.get_messages( + session_id="session-uuid", + edit_strategies=[{"type": "token_limit", "params": {"limit_tokens": 30000}}], + editing_trigger={"token_gte": 50000}, +) print(f"Current tokens: {result.this_time_tokens}") - -if result.this_time_tokens > 50000: - # Apply strategies to reduce context - result = client.sessions.get_messages( - session_id="session-uuid", - edit_strategies=[{"type": "token_limit", "params": {"limit_tokens": 30000}}] - ) ``` ```typescript title="TypeScript" -let result = await client.sessions.getMessages("session-uuid"); +const result = await client.sessions.getMessages("session-uuid", { + editStrategies: [{ type: "token_limit", params: { limit_tokens: 30000 } }], + editingTrigger: { token_gte: 50000 }, +}); console.log(`Current tokens: ${result.thisTimeTokens}`); - -if (result.thisTimeTokens > 50000) { - // Apply strategies to reduce context - result = await client.sessions.getMessages("session-uuid", { - editStrategies: [{ type: "token_limit", params: { limit_tokens: 30000 } }], - }); -} ``` diff --git a/src/client/acontext-py/src/acontext/resources/async_sessions.py b/src/client/acontext-py/src/acontext/resources/async_sessions.py index cfd6a66bf..856cdf61e 100644 --- a/src/client/acontext-py/src/acontext/resources/async_sessions.py +++ b/src/client/acontext-py/src/acontext/resources/async_sessions.py @@ -374,6 +374,8 @@ async def get_messages( format: Literal["acontext", "openai", "anthropic", "gemini"] = "openai", time_desc: bool | None = None, edit_strategies: Optional[List[EditStrategy]] = None, + # editing_trigger triggers edit_strategies (v0 supports {"token_gte": int}). + editing_trigger: dict[str, Any] | None = None, pin_editing_strategies_at_message: str | None = None, ) -> GetMessagesOutput: """Get messages for a session. @@ -394,6 +396,7 @@ async def get_messages( - Middle out: [{"type": "middle_out", "params": {"token_reduce_to": 5000}}] - Token limit: [{"type": "token_limit", "params": {"limit_tokens": 20000}}] Defaults to None. + editing_trigger: Trigger config for edit_strategies, e.g. {"token_gte": 30000}. Defaults to None. pin_editing_strategies_at_message: Message ID to pin editing strategies at. When provided, strategies are only applied to messages up to and including this message ID, keeping subsequent messages unchanged. This helps maintain @@ -419,6 +422,10 @@ async def get_messages( if edit_strategies is not None: validate_edit_strategies(edit_strategies) params["edit_strategies"] = json.dumps(edit_strategies) + if editing_trigger is not None: + if isinstance(editing_trigger, BaseModel): + editing_trigger = editing_trigger.model_dump() + params["editing_trigger"] = json.dumps(editing_trigger) if pin_editing_strategies_at_message is not None: params["pin_editing_strategies_at_message"] = ( pin_editing_strategies_at_message diff --git a/src/client/acontext-py/src/acontext/resources/sessions.py b/src/client/acontext-py/src/acontext/resources/sessions.py index 083504a48..60c19b5ff 100644 --- a/src/client/acontext-py/src/acontext/resources/sessions.py +++ b/src/client/acontext-py/src/acontext/resources/sessions.py @@ -374,6 +374,8 @@ def get_messages( format: Literal["acontext", "openai", "anthropic", "gemini"] = "openai", time_desc: bool | None = None, edit_strategies: Optional[List[EditStrategy]] = None, + # editing_trigger triggers edit_strategies (v0 supports {"token_gte": int}). + editing_trigger: dict[str, Any] | None = None, pin_editing_strategies_at_message: str | None = None, ) -> GetMessagesOutput: """Get messages for a session. @@ -394,6 +396,7 @@ def get_messages( - Middle out: [{"type": "middle_out", "params": {"token_reduce_to": 5000}}] - Token limit: [{"type": "token_limit", "params": {"limit_tokens": 20000}}] Defaults to None. + editing_trigger: Trigger config for edit_strategies, e.g. {"token_gte": 30000}. Defaults to None. pin_editing_strategies_at_message: Message ID to pin editing strategies at. When provided, strategies are only applied to messages up to and including this message ID, keeping subsequent messages unchanged. This helps maintain @@ -419,6 +422,10 @@ def get_messages( if edit_strategies is not None: validate_edit_strategies(edit_strategies) params["edit_strategies"] = json.dumps(edit_strategies) + if editing_trigger is not None: + if isinstance(editing_trigger, BaseModel): + editing_trigger = editing_trigger.model_dump() + params["editing_trigger"] = json.dumps(editing_trigger) if pin_editing_strategies_at_message is not None: params["pin_editing_strategies_at_message"] = ( pin_editing_strategies_at_message diff --git a/src/client/acontext-ts/src/resources/sessions.ts b/src/client/acontext-ts/src/resources/sessions.ts index 2ff2aacdd..6d1cb3460 100644 --- a/src/client/acontext-ts/src/resources/sessions.ts +++ b/src/client/acontext-ts/src/resources/sessions.ts @@ -340,6 +340,7 @@ export class SessionsAPI { * @param options.format - The format of the messages ('acontext', 'openai', 'anthropic', or 'gemini'). * @param options.timeDesc - Order by created_at descending if true, ascending if false. * @param options.editStrategies - Optional list of edit strategies to apply before format conversion. + * @param options.editingTrigger - Optional trigger config for editStrategies (v0 supports { token_gte: number }). * Examples: * - Remove tool results: [{ type: 'remove_tool_result', params: { keep_recent_n_tool_results: 3 } }] * - Remove large tool results: [{ type: 'remove_tool_result', params: { gt_token: 100 } }] @@ -364,6 +365,7 @@ export class SessionsAPI { format?: 'acontext' | 'openai' | 'anthropic' | 'gemini'; timeDesc?: boolean | null; editStrategies?: Array | null; + editingTrigger?: Record | null; pinEditingStrategiesAtMessage?: string | null; } ): Promise { @@ -387,6 +389,9 @@ export class SessionsAPI { EditStrategySchema.array().parse(options.editStrategies); params.edit_strategies = JSON.stringify(options.editStrategies); } + if (options?.editingTrigger !== undefined && options?.editingTrigger !== null) { + params.editing_trigger = JSON.stringify(options.editingTrigger); + } if (options?.pinEditingStrategiesAtMessage !== undefined && options?.pinEditingStrategiesAtMessage !== null) { params.pin_editing_strategies_at_message = options.pinEditingStrategiesAtMessage; } diff --git a/src/server/api/go/internal/modules/handler/session.go b/src/server/api/go/internal/modules/handler/session.go index f8237ecd7..445a848c2 100644 --- a/src/server/api/go/internal/modules/handler/session.go +++ b/src/server/api/go/internal/modules/handler/session.go @@ -495,6 +495,7 @@ type GetMessagesReq struct { Format string `form:"format,default=openai" json:"format" binding:"omitempty,oneof=acontext openai anthropic gemini" example:"openai" enums:"acontext,openai,anthropic,gemini"` TimeDesc bool `form:"time_desc,default=false" json:"time_desc" example:"false"` EditStrategies string `form:"edit_strategies" json:"edit_strategies" example:"[{\"type\":\"remove_tool_result\",\"params\":{\"keep_recent_n_tool_results\":3}}]"` + EditingTrigger string `form:"editing_trigger" json:"editing_trigger" example:"{\"token_gte\":30000}"` PinEditingStrategiesAtMessage string `form:"pin_editing_strategies_at_message" json:"pin_editing_strategies_at_message" example:""` } @@ -513,6 +514,7 @@ type GetMessagesReq struct { // @Param format query string false "Format to convert messages to: acontext (original), openai (default), anthropic, gemini." enums(acontext,openai,anthropic,gemini) // @Param time_desc query boolean false "Order by created_at descending if true, ascending if false (default false)" example(false) // @Param edit_strategies query string false "JSON array of edit strategies to apply before format conversion" example([{"type":"remove_tool_result","params":{"keep_recent_n_tool_results":3}}]) +// @Param editing_trigger query string false "JSON object trigger for edit_strategies. v0 supports only {\"token_gte\": } (OR semantics when more triggers are added)." example({"token_gte":30000}) // @Param pin_editing_strategies_at_message query string false "Message ID to pin editing strategies at. When provided, strategies are only applied to messages up to and including this message ID, keeping subsequent messages unchanged. This helps maintain prompt cache stability by preserving a stable prefix. The response will include edit_at_message_id indicating where strategies were applied." example() // @Security BearerAuth // @Success 200 {object} serializer.Response{data=converter.GetMessagesOutput} @@ -552,6 +554,58 @@ func (h *SessionHandler) GetMessages(c *gin.Context) { } } + // Parse editing_trigger if provided (v0 supports only token_gte). + var editingTrigger *service.EditingTrigger + if req.EditingTrigger != "" { + if req.EditStrategies == "" { + c.JSON(http.StatusBadRequest, serializer.ParamErr("editing_trigger requires edit_strategies", errors.New("missing edit_strategies"))) + return + } + + var raw map[string]interface{} + if err := sonic.Unmarshal([]byte(req.EditingTrigger), &raw); err != nil { + c.JSON(http.StatusBadRequest, serializer.ParamErr("invalid editing_trigger JSON", err)) + return + } + allowedTriggerKeys := map[string]struct{}{ + "token_gte": {}, + } + for k := range raw { + if _, ok := allowedTriggerKeys[k]; !ok { + c.JSON(http.StatusBadRequest, serializer.ParamErr("invalid editing_trigger", fmt.Errorf("unsupported trigger: %s", k))) + return + } + } + + var trig service.EditingTrigger + if err := sonic.Unmarshal([]byte(req.EditingTrigger), &trig); err != nil { + c.JSON(http.StatusBadRequest, serializer.ParamErr("invalid editing_trigger JSON", err)) + return + } + + triggerValidators := map[string]func(service.EditingTrigger) error{ + "token_gte": func(t service.EditingTrigger) error { + if t.TokenGte == nil || *t.TokenGte <= 0 { + return errors.New("token_gte must be > 0") + } + return nil + }, + } + for k := range raw { + validate, ok := triggerValidators[k] + if !ok { + // Should be unreachable due to allowedTriggerKeys check above. + c.JSON(http.StatusBadRequest, serializer.ParamErr("invalid editing_trigger", fmt.Errorf("unsupported trigger: %s", k))) + return + } + if err := validate(trig); err != nil { + c.JSON(http.StatusBadRequest, serializer.ParamErr("invalid editing_trigger."+k, err)) + return + } + } + editingTrigger = &trig + } + out, err := h.svc.GetMessages(c.Request.Context(), service.GetMessagesInput{ ProjectID: project.ID, SessionID: sessionID, @@ -562,6 +616,7 @@ func (h *SessionHandler) GetMessages(c *gin.Context) { AssetExpire: time.Hour * 24, TimeDesc: req.TimeDesc, EditStrategies: editStrategies, + EditingTrigger: editingTrigger, PinEditingStrategiesAtMessage: req.PinEditingStrategiesAtMessage, UserKEK: middleware.GetUserKEKIfEncrypted(c), }) diff --git a/src/server/api/go/internal/modules/service/session.go b/src/server/api/go/internal/modules/service/session.go index d1fa0a33f..343b1a8a7 100644 --- a/src/server/api/go/internal/modules/service/session.go +++ b/src/server/api/go/internal/modules/service/session.go @@ -21,6 +21,7 @@ import ( "github.com/memodb-io/Acontext/internal/modules/repo" "github.com/memodb-io/Acontext/internal/pkg/editor" "github.com/memodb-io/Acontext/internal/pkg/paging" + "github.com/memodb-io/Acontext/internal/pkg/tokenizer" "github.com/redis/go-redis/v9" "go.uber.org/zap" "gorm.io/datatypes" @@ -436,17 +437,69 @@ func (s *sessionService) StoreMessage(ctx context.Context, in StoreMessageInput) } type GetMessagesInput struct { - ProjectID uuid.UUID `json:"project_id"` - SessionID uuid.UUID `json:"session_id"` - Limit int `json:"limit"` - Cursor string `json:"cursor"` - WithAssetPublicURL bool `json:"with_public_url"` - AssetExpire time.Duration `json:"asset_expire"` - TimeDesc bool `json:"time_desc"` - WithEvents bool `json:"with_events"` - EditStrategies []editor.StrategyConfig `json:"edit_strategies,omitempty"` - PinEditingStrategiesAtMessage string `json:"pin_editing_strategies_at_message,omitempty"` - UserKEK []byte `json:"-"` // optional: for envelope encryption (decrypting parts) + ProjectID uuid.UUID `json:"project_id"` + SessionID uuid.UUID `json:"session_id"` + Limit int `json:"limit"` + Cursor string `json:"cursor"` + WithAssetPublicURL bool `json:"with_public_url"` + AssetExpire time.Duration `json:"asset_expire"` + TimeDesc bool `json:"time_desc"` + WithEvents bool `json:"with_events"` + EditStrategies []editor.StrategyConfig `json:"edit_strategies,omitempty"` + // EditingTrigger holds optional trigger config for applying edit_strategies. + EditingTrigger *EditingTrigger `json:"editing_trigger,omitempty"` + PinEditingStrategiesAtMessage string `json:"pin_editing_strategies_at_message,omitempty"` + UserKEK []byte `json:"-"` // optional: for envelope encryption (decrypting parts) +} + +// EditingTrigger defines trigger configuration for applying edit_strategies. +// v0 supports only token_gte. +type EditingTrigger struct { + // TokenGte triggers edit strategies when the token count is greater than or equal to this value. + TokenGte *int `json:"token_gte,omitempty"` +} + +type editingTriggerEval struct { + sessionID uuid.UUID + messages []model.Message + + tokenCount *int +} + +func (e *editingTriggerEval) Tokens(ctx context.Context) (int, error) { + if e.tokenCount != nil { + return *e.tokenCount, nil + } + + tokens, err := tokenizer.CountMessagePartsTokens(ctx, e.messages) + if err != nil { + return 0, fmt.Errorf("failed to count tokens for editing_trigger session_id=%s: %w", e.sessionID, err) + } + + e.tokenCount = &tokens + return tokens, nil +} + +type editingTriggerCheck func(ctx context.Context, eval *editingTriggerEval) (bool, error) + +func buildEditingTriggerChecks(trigger *EditingTrigger) []editingTriggerCheck { + checks := make([]editingTriggerCheck, 0, 1) + if trigger == nil { + return checks + } + + if trigger.TokenGte != nil && *trigger.TokenGte > 0 { + threshold := *trigger.TokenGte + checks = append(checks, func(ctx context.Context, eval *editingTriggerEval) (bool, error) { + tokens, err := eval.Tokens(ctx) + if err != nil { + return false, err + } + return tokens >= threshold, nil + }) + } + + return checks } type PublicURL struct { @@ -552,12 +605,59 @@ func (s *sessionService) GetMessages(ctx context.Context, in GetMessagesInput) ( // Apply edit strategies if provided (before format conversion) if len(in.EditStrategies) > 0 { - result, err := editor.ApplyStrategiesWithPin(out.Items, in.EditStrategies, in.PinEditingStrategiesAtMessage) - if err != nil { - return nil, fmt.Errorf("failed to apply edit strategies: %w", err) + applyEditStrategies := true + triggerChecks := buildEditingTriggerChecks(in.EditingTrigger) + if len(triggerChecks) > 0 { + // Evaluate trigger on the same editable prefix used by pin_editing_strategies_at_message. + triggerMessages := out.Items + effectivePin := "" + if in.PinEditingStrategiesAtMessage != "" { + pinIndex := -1 + for i := range out.Items { + if out.Items[i].ID.String() == in.PinEditingStrategiesAtMessage { + pinIndex = i + break + } + } + if pinIndex != -1 { + effectivePin = in.PinEditingStrategiesAtMessage + triggerMessages = out.Items[:pinIndex+1] + } + } + + // OR semantics: apply when any trigger check passes. + applyEditStrategies = false + eval := &editingTriggerEval{sessionID: in.SessionID, messages: triggerMessages} + for _, check := range triggerChecks { + ok, err := check(ctx, eval) + if err != nil { + return nil, err + } + if ok { + applyEditStrategies = true + break + } + } + + if !applyEditStrategies && len(out.Items) > 0 { + if effectivePin != "" { + out.EditAtMessageID = effectivePin + } else { + out.EditAtMessageID = out.Items[len(out.Items)-1].ID.String() + } + } + } + + if applyEditStrategies { + result, err := editor.ApplyStrategiesWithPin(out.Items, in.EditStrategies, in.PinEditingStrategiesAtMessage) + if err != nil { + return nil, fmt.Errorf("failed to apply edit strategies: %w", err) + } + out.Items = result.Messages + out.EditAtMessageID = result.EditAtMessageID + } else if out.EditAtMessageID == "" && len(out.Items) > 0 { + out.EditAtMessageID = out.Items[len(out.Items)-1].ID.String() } - out.Items = result.Messages - out.EditAtMessageID = result.EditAtMessageID } else if len(out.Items) > 0 { // No strategies, but still set EditAtMessageID to the last message out.EditAtMessageID = out.Items[len(out.Items)-1].ID.String() From d47fb9538a92814351e9ca1c9989ad00645ebd2c Mon Sep 17 00:00:00 2001 From: riturajFi Date: Fri, 13 Feb 2026 01:47:30 +0530 Subject: [PATCH 02/20] temporary removal of edited at --- .../api/go/internal/modules/service/session.go | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/src/server/api/go/internal/modules/service/session.go b/src/server/api/go/internal/modules/service/session.go index 343b1a8a7..ebda57a05 100644 --- a/src/server/api/go/internal/modules/service/session.go +++ b/src/server/api/go/internal/modules/service/session.go @@ -610,7 +610,6 @@ func (s *sessionService) GetMessages(ctx context.Context, in GetMessagesInput) ( if len(triggerChecks) > 0 { // Evaluate trigger on the same editable prefix used by pin_editing_strategies_at_message. triggerMessages := out.Items - effectivePin := "" if in.PinEditingStrategiesAtMessage != "" { pinIndex := -1 for i := range out.Items { @@ -620,7 +619,6 @@ func (s *sessionService) GetMessages(ctx context.Context, in GetMessagesInput) ( } } if pinIndex != -1 { - effectivePin = in.PinEditingStrategiesAtMessage triggerMessages = out.Items[:pinIndex+1] } } @@ -639,13 +637,15 @@ func (s *sessionService) GetMessages(ctx context.Context, in GetMessagesInput) ( } } - if !applyEditStrategies && len(out.Items) > 0 { - if effectivePin != "" { - out.EditAtMessageID = effectivePin - } else { - out.EditAtMessageID = out.Items[len(out.Items)-1].ID.String() - } - } + // TODO(maintainers): clarify final semantics for edit_at_message_id when trigger does not fire. + // Should it remain empty (strict "strategies were applied" meaning), or keep a reusable pin hint? + // if !applyEditStrategies && len(out.Items) > 0 { + // if effectivePin != "" { + // out.EditAtMessageID = effectivePin + // } else { + // out.EditAtMessageID = out.Items[len(out.Items)-1].ID.String() + // } + // } } if applyEditStrategies { From d80158b8676ec03735ed2f8b445f9c99a88e476e Mon Sep 17 00:00:00 2001 From: riturajFi Date: Fri, 13 Feb 2026 15:55:18 +0530 Subject: [PATCH 03/20] minor fix --- .../go/internal/modules/handler/session.go | 13 ++-- .../internal/modules/handler/session_test.go | 64 +++++++++++++++++ .../go/internal/modules/service/session.go | 27 ++++++++ .../internal/modules/service/session_test.go | 69 ++++++++++++++++++- 4 files changed, 163 insertions(+), 10 deletions(-) diff --git a/src/server/api/go/internal/modules/handler/session.go b/src/server/api/go/internal/modules/handler/session.go index 445a848c2..fca08b5d0 100644 --- a/src/server/api/go/internal/modules/handler/session.go +++ b/src/server/api/go/internal/modules/handler/session.go @@ -621,6 +621,10 @@ func (h *SessionHandler) GetMessages(c *gin.Context) { UserKEK: middleware.GetUserKEKIfEncrypted(c), }) if err != nil { + if errors.Is(err, service.ErrGetMessagesTokenCount) { + c.JSON(http.StatusInternalServerError, serializer.DBErr("failed to count tokens", err)) + return + } c.JSON(http.StatusBadRequest, serializer.DBErr("", err)) return } @@ -637,13 +641,6 @@ func (h *SessionHandler) GetMessages(c *gin.Context) { return } - // Calculate token count for the returned messages - thisTimeTokens, err := tokenizer.CountMessagePartsTokens(c.Request.Context(), out.Items) - if err != nil { - c.JSON(http.StatusInternalServerError, serializer.DBErr("failed to count tokens", err)) - return - } - convertedOut, err := converter.GetConvertedMessagesOutput( out.Items, format, @@ -651,7 +648,7 @@ func (h *SessionHandler) GetMessages(c *gin.Context) { out.Events, out.NextCursor, out.HasMore, - thisTimeTokens, + out.ThisTimeTokens, out.EditAtMessageID, ) if err != nil { diff --git a/src/server/api/go/internal/modules/handler/session_test.go b/src/server/api/go/internal/modules/handler/session_test.go index 42b24b929..1c65ba272 100644 --- a/src/server/api/go/internal/modules/handler/session_test.go +++ b/src/server/api/go/internal/modules/handler/session_test.go @@ -4288,3 +4288,67 @@ func TestSessionHandler_DownloadSessionAsset(t *testing.T) { mockService.AssertExpectations(t) }) } + +func TestSessionHandler_GetMessages_ServiceTokenCountError(t *testing.T) { + gin.SetMode(gin.TestMode) + + projectID := uuid.New() + sessionID := uuid.New() + mockService := &MockSessionService{} + mockService.On("GetMessages", mock.Anything, mock.Anything).Return(nil, service.ErrGetMessagesTokenCount) + + handler := NewSessionHandler(mockService, &MockUserService{}, getMockSessionCoreClient()) + router := setupSessionRouter() + router.GET("/session/:session_id/messages", func(c *gin.Context) { + c.Set("project", &model.Project{ID: projectID}) + handler.GetMessages(c) + }) + + req := httptest.NewRequest("GET", "/session/"+sessionID.String()+"/messages?limit=20", nil) + w := httptest.NewRecorder() + router.ServeHTTP(w, req) + + assert.Equal(t, http.StatusInternalServerError, w.Code) + mockService.AssertExpectations(t) +} + +func TestSessionHandler_GetMessages_UsesServiceThisTimeTokens(t *testing.T) { + gin.SetMode(gin.TestMode) + + projectID := uuid.New() + sessionID := uuid.New() + mockService := &MockSessionService{} + mockService.On("GetMessages", mock.Anything, mock.Anything).Return(&service.GetMessagesOutput{ + Items: []model.Message{ + { + ID: uuid.New(), + SessionID: sessionID, + Role: model.RoleUser, + }, + }, + HasMore: false, + ThisTimeTokens: 999, + }, nil) + + handler := NewSessionHandler(mockService, &MockUserService{}, getMockSessionCoreClient()) + router := setupSessionRouter() + router.GET("/session/:session_id/messages", func(c *gin.Context) { + c.Set("project", &model.Project{ID: projectID}) + handler.GetMessages(c) + }) + + req := httptest.NewRequest("GET", "/session/"+sessionID.String()+"/messages?limit=20", nil) + w := httptest.NewRecorder() + router.ServeHTTP(w, req) + + assert.Equal(t, http.StatusOK, w.Code) + mockService.AssertExpectations(t) + + var response map[string]interface{} + err := sonic.Unmarshal(w.Body.Bytes(), &response) + require.NoError(t, err) + + data, ok := response["data"].(map[string]interface{}) + require.True(t, ok) + assert.Equal(t, float64(999), data["this_time_tokens"]) +} diff --git a/src/server/api/go/internal/modules/service/session.go b/src/server/api/go/internal/modules/service/session.go index ebda57a05..2c07fb9da 100644 --- a/src/server/api/go/internal/modules/service/session.go +++ b/src/server/api/go/internal/modules/service/session.go @@ -68,6 +68,8 @@ type sessionService struct { materialSvc MaterialService } +var ErrGetMessagesTokenCount = errors.New("get messages token count error") + const ( // Redis key prefix for message parts cache redisKeyPrefixParts = "message:parts:" @@ -502,6 +504,18 @@ func buildEditingTriggerChecks(trigger *EditingTrigger) []editingTriggerCheck { return checks } +func sameMessageOrderByID(a, b []model.Message) bool { + if len(a) != len(b) { + return false + } + for i := range a { + if a[i].ID != b[i].ID { + return false + } + } + return true +} + type PublicURL struct { URL string `json:"url"` ExpireAt time.Time `json:"expire_at"` @@ -513,6 +527,7 @@ type GetMessagesOutput struct { NextCursor string `json:"next_cursor,omitempty"` HasMore bool `json:"has_more"` PublicURLs map[string]PublicURL `json:"public_urls,omitempty"` // file_name -> url + ThisTimeTokens int `json:"this_time_tokens"` EditAtMessageID string `json:"edit_at_message_id,omitempty"` } @@ -604,6 +619,7 @@ func (s *sessionService) GetMessages(ctx context.Context, in GetMessagesInput) ( } // Apply edit strategies if provided (before format conversion) + var triggerEval *editingTriggerEval if len(in.EditStrategies) > 0 { applyEditStrategies := true triggerChecks := buildEditingTriggerChecks(in.EditingTrigger) @@ -626,6 +642,7 @@ func (s *sessionService) GetMessages(ctx context.Context, in GetMessagesInput) ( // OR semantics: apply when any trigger check passes. applyEditStrategies = false eval := &editingTriggerEval{sessionID: in.SessionID, messages: triggerMessages} + triggerEval = eval for _, check := range triggerChecks { ok, err := check(ctx, eval) if err != nil { @@ -688,6 +705,16 @@ func (s *sessionService) GetMessages(ctx context.Context, in GetMessagesInput) ( } } + if triggerEval != nil && triggerEval.tokenCount != nil && sameMessageOrderByID(triggerEval.messages, out.Items) { + out.ThisTimeTokens = *triggerEval.tokenCount + } else { + thisTimeTokens, err := tokenizer.CountMessagePartsTokens(ctx, out.Items) + if err != nil { + return nil, fmt.Errorf("%w: session_id=%s: %v", ErrGetMessagesTokenCount, in.SessionID, err) + } + out.ThisTimeTokens = thisTimeTokens + } + return out, nil } diff --git a/src/server/api/go/internal/modules/service/session_test.go b/src/server/api/go/internal/modules/service/session_test.go index b79e5175f..9898fa185 100644 --- a/src/server/api/go/internal/modules/service/session_test.go +++ b/src/server/api/go/internal/modules/service/session_test.go @@ -13,6 +13,7 @@ import ( "github.com/memodb-io/Acontext/internal/config" "github.com/memodb-io/Acontext/internal/modules/model" "github.com/memodb-io/Acontext/internal/modules/repo" + "github.com/memodb-io/Acontext/internal/pkg/tokenizer" "github.com/redis/go-redis/v9" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/mock" @@ -1501,6 +1502,9 @@ func TestSessionService_GetMessages_MaterialURLs(t *testing.T) { projectID := uuid.New() sessionID := uuid.New() + err := tokenizer.Init(logger) + assert.NoError(t, err) + // Parts to cache in Redis so loadPartsForMessage succeeds without S3 imageParts := []model.Part{ { @@ -1570,7 +1574,9 @@ func TestSessionService_GetMessages_MaterialURLs(t *testing.T) { AssetExpire: time.Hour, }) - assert.NoError(t, err) + if !assert.NoError(t, err) { + return + } assert.NotNil(t, result) assert.NotEmpty(t, result.PublicURLs) @@ -1617,10 +1623,69 @@ func TestSessionService_GetMessages_MaterialURLs(t *testing.T) { WithAssetPublicURL: false, }) - assert.NoError(t, err) + if !assert.NoError(t, err) { + return + } assert.Empty(t, result.PublicURLs) // materialSvc should NOT have been called mockMaterialSvc.AssertNotCalled(t, "CreateMaterialURL") }) } + +func TestSessionService_GetMessages_ComputesThisTimeTokens(t *testing.T) { + ctx := context.Background() + projectID := uuid.New() + sessionID := uuid.New() + logger := zap.NewNop() + cfg := &config.Config{} + + err := tokenizer.Init(logger) + assert.NoError(t, err) + + mr := miniredis.RunT(t) + rdb := redis.NewClient(&redis.Options{Addr: mr.Addr()}) + + repo := &MockSessionRepo{} + session := &model.Session{ID: sessionID, ProjectID: projectID} + repo.On("Get", mock.Anything, mock.MatchedBy(func(s *model.Session) bool { + return s.ID == sessionID + })).Return(session, nil) + repo.On("ListAllMessagesBySession", ctx, sessionID).Return([]model.Message{ + { + ID: uuid.New(), + SessionID: sessionID, + Role: model.RoleUser, + CreatedAt: time.Now(), + PartsAssetMeta: datatypes.NewJSONType(model.Asset{ + S3Key: "parts/" + projectID.String() + "/message.json", + SHA256: "sha-message", + MIME: "application/json", + }), + }, + }, nil) + + partsJSON, err := json.Marshal([]model.Part{ + model.NewTextPart("hello from acontext"), + }) + assert.NoError(t, err) + err = rdb.Set(ctx, "message:parts:"+projectID.String()+":sha-message", append([]byte{0x00}, partsJSON...), time.Hour).Err() + assert.NoError(t, err) + + mockAssetRefRepo := &MockAssetReferenceRepo{} + service := NewSessionService(repo, nil, mockAssetRefRepo, nil, logger, nil, nil, cfg, rdb, nil) + + out, err := service.GetMessages(ctx, GetMessagesInput{ + ProjectID: projectID, + SessionID: sessionID, + Limit: 0, + }) + assert.NoError(t, err) + assert.NotNil(t, out) + assert.Greater(t, out.ThisTimeTokens, 0) + assert.Len(t, out.Items, 1) + assert.Len(t, out.Items[0].Parts, 1) + assert.Equal(t, "hello from acontext", out.Items[0].Parts[0].Text) + + repo.AssertExpectations(t) +} From 3985a960d3fb23a9effe012324987e2fa34de911 Mon Sep 17 00:00:00 2001 From: riturajFi Date: Sat, 14 Feb 2026 01:16:22 +0530 Subject: [PATCH 04/20] minor fix --- .../go/internal/modules/service/session.go | 6 +- .../internal/modules/service/session_test.go | 106 ++++++++++++++++++ 2 files changed, 110 insertions(+), 2 deletions(-) diff --git a/src/server/api/go/internal/modules/service/session.go b/src/server/api/go/internal/modules/service/session.go index 2c07fb9da..05dc9d7a6 100644 --- a/src/server/api/go/internal/modules/service/session.go +++ b/src/server/api/go/internal/modules/service/session.go @@ -475,7 +475,7 @@ func (e *editingTriggerEval) Tokens(ctx context.Context) (int, error) { tokens, err := tokenizer.CountMessagePartsTokens(ctx, e.messages) if err != nil { - return 0, fmt.Errorf("failed to count tokens for editing_trigger session_id=%s: %w", e.sessionID, err) + return 0, fmt.Errorf("%w: failed to count tokens for editing_trigger session_id=%s: %v", ErrGetMessagesTokenCount, e.sessionID, err) } e.tokenCount = &tokens @@ -620,6 +620,7 @@ func (s *sessionService) GetMessages(ctx context.Context, in GetMessagesInput) ( // Apply edit strategies if provided (before format conversion) var triggerEval *editingTriggerEval + strategiesApplied := false if len(in.EditStrategies) > 0 { applyEditStrategies := true triggerChecks := buildEditingTriggerChecks(in.EditingTrigger) @@ -670,6 +671,7 @@ func (s *sessionService) GetMessages(ctx context.Context, in GetMessagesInput) ( if err != nil { return nil, fmt.Errorf("failed to apply edit strategies: %w", err) } + strategiesApplied = true out.Items = result.Messages out.EditAtMessageID = result.EditAtMessageID } else if out.EditAtMessageID == "" && len(out.Items) > 0 { @@ -705,7 +707,7 @@ func (s *sessionService) GetMessages(ctx context.Context, in GetMessagesInput) ( } } - if triggerEval != nil && triggerEval.tokenCount != nil && sameMessageOrderByID(triggerEval.messages, out.Items) { + if triggerEval != nil && triggerEval.tokenCount != nil && !strategiesApplied && sameMessageOrderByID(triggerEval.messages, out.Items) { out.ThisTimeTokens = *triggerEval.tokenCount } else { thisTimeTokens, err := tokenizer.CountMessagePartsTokens(ctx, out.Items) diff --git a/src/server/api/go/internal/modules/service/session_test.go b/src/server/api/go/internal/modules/service/session_test.go index 9898fa185..6d88fb3cc 100644 --- a/src/server/api/go/internal/modules/service/session_test.go +++ b/src/server/api/go/internal/modules/service/session_test.go @@ -5,6 +5,7 @@ import ( "encoding/json" "errors" "fmt" + "strings" "testing" "time" @@ -13,6 +14,7 @@ import ( "github.com/memodb-io/Acontext/internal/config" "github.com/memodb-io/Acontext/internal/modules/model" "github.com/memodb-io/Acontext/internal/modules/repo" + "github.com/memodb-io/Acontext/internal/pkg/editor" "github.com/memodb-io/Acontext/internal/pkg/tokenizer" "github.com/redis/go-redis/v9" "github.com/stretchr/testify/assert" @@ -1689,3 +1691,107 @@ func TestSessionService_GetMessages_ComputesThisTimeTokens(t *testing.T) { repo.AssertExpectations(t) } + +func TestSessionService_GetMessages_ThisTimeTokensMatchesEditedOutput(t *testing.T) { + ctx := context.Background() + sessionID := uuid.New() + triggerThreshold := 1 + + err := tokenizer.Init(zap.NewNop()) + assert.NoError(t, err) + + repo := &MockSessionRepo{} + repoMessages := []model.Message{ + { + ID: uuid.New(), + SessionID: sessionID, + Role: model.RoleAssistant, + Parts: []model.Part{ + { + Type: model.PartTypeToolResult, + Text: strings.Repeat("very large tool result payload ", 200), + }, + }, + }, + } + repo.On("ListAllMessagesBySession", ctx, sessionID).Return(repoMessages, nil) + + preEditTokens, err := tokenizer.CountMessagePartsTokens(ctx, repoMessages) + assert.NoError(t, err) + assert.Greater(t, preEditTokens, 0) + + service := NewSessionService(repo, &MockAssetReferenceRepo{}, zap.NewNop(), nil, nil, &config.Config{}, nil) + out, err := service.GetMessages(ctx, GetMessagesInput{ + SessionID: sessionID, + Limit: 0, + EditStrategies: []editor.StrategyConfig{ + { + Type: "remove_tool_result", + Params: map[string]interface{}{ + "keep_recent_n_tool_results": 0, + }, + }, + }, + EditingTrigger: &EditingTrigger{ + TokenGte: &triggerThreshold, + }, + }) + assert.NoError(t, err) + assert.NotNil(t, out) + assert.Len(t, out.Items, 1) + assert.Equal(t, "Done", out.Items[0].Parts[0].Text) + + finalTokens, err := tokenizer.CountMessagePartsTokens(ctx, out.Items) + assert.NoError(t, err) + assert.Equal(t, finalTokens, out.ThisTimeTokens) + assert.Less(t, out.ThisTimeTokens, preEditTokens) + + repo.AssertExpectations(t) +} + +func TestSessionService_GetMessages_TriggerTokenErrorsAreWrapped(t *testing.T) { + ctx := context.Background() + sessionID := uuid.New() + triggerThreshold := 1 + + repo := &MockSessionRepo{} + repo.On("ListAllMessagesBySession", ctx, sessionID).Return([]model.Message{ + { + ID: uuid.New(), + SessionID: sessionID, + Role: model.RoleAssistant, + Parts: []model.Part{ + { + Type: model.PartTypeToolCall, + Meta: map[string]interface{}{ + model.MetaKeyName: "bad_tool", + model.MetaKeyArguments: func() {}, + }, + }, + }, + }, + }, nil) + + service := NewSessionService(repo, &MockAssetReferenceRepo{}, zap.NewNop(), nil, nil, &config.Config{}, nil) + out, err := service.GetMessages(ctx, GetMessagesInput{ + SessionID: sessionID, + Limit: 0, + EditStrategies: []editor.StrategyConfig{ + { + Type: "token_limit", + Params: map[string]interface{}{ + "limit_tokens": 10, + }, + }, + }, + EditingTrigger: &EditingTrigger{ + TokenGte: &triggerThreshold, + }, + }) + + assert.Nil(t, out) + assert.Error(t, err) + assert.ErrorIs(t, err, ErrGetMessagesTokenCount) + + repo.AssertExpectations(t) +} From ef1b0f73202414d2ada0cfbe0db2c7f5fa420e01 Mon Sep 17 00:00:00 2001 From: riturajFi Date: Sun, 15 Feb 2026 18:37:40 +0530 Subject: [PATCH 05/20] temp --- .env.local-api | 30 +++++++++++++++++++ src/server/.env.local-api | 30 +++++++++++++++++++ .../go/internal/modules/handler/session.go | 4 +++ .../internal/modules/handler/session_test.go | 22 ++++++++++++++ 4 files changed, 86 insertions(+) create mode 100644 .env.local-api create mode 100644 src/server/.env.local-api diff --git a/.env.local-api b/.env.local-api new file mode 100644 index 000000000..05ea7dab6 --- /dev/null +++ b/.env.local-api @@ -0,0 +1,30 @@ +APP_ENV=debug +API_EXPORT_PORT=8029 +ROOT_API_BEARER_TOKEN=local-dev-root + +DATABASE_HOST=127.0.0.1 +DATABASE_EXPORT_PORT=15432 +DATABASE_USER=acontext +DATABASE_PASSWORD=helloworld +DATABASE_NAME=acontext + +REDIS_HOST=127.0.0.1 +REDIS_EXPORT_PORT=16379 +REDIS_PASSWORD=helloworld + +RABBITMQ_HOST=127.0.0.1 +RABBITMQ_EXPORT_PORT=15672 +RABBITMQ_USER=acontext +RABBITMQ_PASSWORD=helloworld +RABBITMQ_VHOST=/ +RABBITMQ_VHOST_ENCODED=%2F + +S3_ENDPOINT=http://127.0.0.1:19000 +S3_INTERNAL_ENDPOINT=http://127.0.0.1:19000 +S3_REGION=auto +S3_ACCESS_KEY=acontext +S3_SECRET_KEY=helloworld +S3_BUCKET=acontext-assets + +CORE_BASE_URL=http://127.0.0.1:8019 +OTEL_EXPORTER_OTLP_ENDPOINT= diff --git a/src/server/.env.local-api b/src/server/.env.local-api new file mode 100644 index 000000000..05ea7dab6 --- /dev/null +++ b/src/server/.env.local-api @@ -0,0 +1,30 @@ +APP_ENV=debug +API_EXPORT_PORT=8029 +ROOT_API_BEARER_TOKEN=local-dev-root + +DATABASE_HOST=127.0.0.1 +DATABASE_EXPORT_PORT=15432 +DATABASE_USER=acontext +DATABASE_PASSWORD=helloworld +DATABASE_NAME=acontext + +REDIS_HOST=127.0.0.1 +REDIS_EXPORT_PORT=16379 +REDIS_PASSWORD=helloworld + +RABBITMQ_HOST=127.0.0.1 +RABBITMQ_EXPORT_PORT=15672 +RABBITMQ_USER=acontext +RABBITMQ_PASSWORD=helloworld +RABBITMQ_VHOST=/ +RABBITMQ_VHOST_ENCODED=%2F + +S3_ENDPOINT=http://127.0.0.1:19000 +S3_INTERNAL_ENDPOINT=http://127.0.0.1:19000 +S3_REGION=auto +S3_ACCESS_KEY=acontext +S3_SECRET_KEY=helloworld +S3_BUCKET=acontext-assets + +CORE_BASE_URL=http://127.0.0.1:8019 +OTEL_EXPORTER_OTLP_ENDPOINT= diff --git a/src/server/api/go/internal/modules/handler/session.go b/src/server/api/go/internal/modules/handler/session.go index fca08b5d0..46812c605 100644 --- a/src/server/api/go/internal/modules/handler/session.go +++ b/src/server/api/go/internal/modules/handler/session.go @@ -567,6 +567,10 @@ func (h *SessionHandler) GetMessages(c *gin.Context) { c.JSON(http.StatusBadRequest, serializer.ParamErr("invalid editing_trigger JSON", err)) return } + if len(raw) == 0 { + c.JSON(http.StatusBadRequest, serializer.ParamErr("invalid editing_trigger", errors.New("at least one supported trigger is required"))) + return + } allowedTriggerKeys := map[string]struct{}{ "token_gte": {}, } diff --git a/src/server/api/go/internal/modules/handler/session_test.go b/src/server/api/go/internal/modules/handler/session_test.go index 1c65ba272..afcfb331e 100644 --- a/src/server/api/go/internal/modules/handler/session_test.go +++ b/src/server/api/go/internal/modules/handler/session_test.go @@ -7,6 +7,7 @@ import ( "mime/multipart" "net/http" "net/http/httptest" + "net/url" "testing" "github.com/bytedance/sonic" @@ -4352,3 +4353,24 @@ func TestSessionHandler_GetMessages_UsesServiceThisTimeTokens(t *testing.T) { require.True(t, ok) assert.Equal(t, float64(999), data["this_time_tokens"]) } + +func TestSessionHandler_GetMessages_RejectsEmptyEditingTrigger(t *testing.T) { + gin.SetMode(gin.TestMode) + + sessionID := uuid.New() + mockService := &MockSessionService{} + + handler := NewSessionHandler(mockService, &MockUserService{}, getMockSessionCoreClient()) + router := setupSessionRouter() + router.GET("/session/:session_id/messages", handler.GetMessages) + + editStrategies := `[{"type":"token_limit","params":{"limit_tokens":100}}]` + reqURL := "/session/" + sessionID.String() + "/messages?limit=20&edit_strategies=" + + url.QueryEscape(editStrategies) + "&editing_trigger=" + url.QueryEscape(`{}`) + req := httptest.NewRequest("GET", reqURL, nil) + w := httptest.NewRecorder() + router.ServeHTTP(w, req) + + assert.Equal(t, http.StatusBadRequest, w.Code) + mockService.AssertNotCalled(t, "GetMessages") +} From 97c34105bea3e0d567d33950e13dcc02ff9ef4ba Mon Sep 17 00:00:00 2001 From: riturajFi Date: Sat, 21 Feb 2026 18:56:10 +0530 Subject: [PATCH 06/20] minor fixes --- .../go/internal/modules/service/session.go | 15 +- .../internal/modules/service/session_test.go | 451 ++++++++++++++++++ 2 files changed, 456 insertions(+), 10 deletions(-) diff --git a/src/server/api/go/internal/modules/service/session.go b/src/server/api/go/internal/modules/service/session.go index 05dc9d7a6..73b5f3f38 100644 --- a/src/server/api/go/internal/modules/service/session.go +++ b/src/server/api/go/internal/modules/service/session.go @@ -624,7 +624,8 @@ func (s *sessionService) GetMessages(ctx context.Context, in GetMessagesInput) ( if len(in.EditStrategies) > 0 { applyEditStrategies := true triggerChecks := buildEditingTriggerChecks(in.EditingTrigger) - if len(triggerChecks) > 0 { + triggerEvaluated := len(triggerChecks) > 0 + if triggerEvaluated { // Evaluate trigger on the same editable prefix used by pin_editing_strategies_at_message. triggerMessages := out.Items if in.PinEditingStrategiesAtMessage != "" { @@ -655,15 +656,6 @@ func (s *sessionService) GetMessages(ctx context.Context, in GetMessagesInput) ( } } - // TODO(maintainers): clarify final semantics for edit_at_message_id when trigger does not fire. - // Should it remain empty (strict "strategies were applied" meaning), or keep a reusable pin hint? - // if !applyEditStrategies && len(out.Items) > 0 { - // if effectivePin != "" { - // out.EditAtMessageID = effectivePin - // } else { - // out.EditAtMessageID = out.Items[len(out.Items)-1].ID.String() - // } - // } } if applyEditStrategies { @@ -674,6 +666,9 @@ func (s *sessionService) GetMessages(ctx context.Context, in GetMessagesInput) ( strategiesApplied = true out.Items = result.Messages out.EditAtMessageID = result.EditAtMessageID + } else if triggerEvaluated && in.PinEditingStrategiesAtMessage != "" { + // Trigger skipped editing; preserve caller-provided boundary for future requests. + out.EditAtMessageID = in.PinEditingStrategiesAtMessage } else if out.EditAtMessageID == "" && len(out.Items) > 0 { out.EditAtMessageID = out.Items[len(out.Items)-1].ID.String() } diff --git a/src/server/api/go/internal/modules/service/session_test.go b/src/server/api/go/internal/modules/service/session_test.go index 6d88fb3cc..b8ad0b7d3 100644 --- a/src/server/api/go/internal/modules/service/session_test.go +++ b/src/server/api/go/internal/modules/service/session_test.go @@ -1795,3 +1795,454 @@ func TestSessionService_GetMessages_TriggerTokenErrorsAreWrapped(t *testing.T) { repo.AssertExpectations(t) } + +func TestSessionService_GetMessages_TriggerFalseKeepsProvidedPin(t *testing.T) { + ctx := context.Background() + sessionID := uuid.New() + triggerThreshold := 100000 + + err := tokenizer.Init(zap.NewNop()) + assert.NoError(t, err) + + pinMsgID := uuid.New() + lastMsgID := uuid.New() + repoMessages := []model.Message{ + { + ID: pinMsgID, + SessionID: sessionID, + Role: model.RoleUser, + Parts: []model.Part{ + model.NewTextPart("short message"), + }, + }, + { + ID: lastMsgID, + SessionID: sessionID, + Role: model.RoleUser, + Parts: []model.Part{ + model.NewTextPart("another short message"), + }, + }, + } + + repo := &MockSessionRepo{} + repo.On("ListAllMessagesBySession", ctx, sessionID).Return(repoMessages, nil) + + service := NewSessionService(repo, &MockAssetReferenceRepo{}, zap.NewNop(), nil, nil, &config.Config{}, nil) + out, err := service.GetMessages(ctx, GetMessagesInput{ + SessionID: sessionID, + Limit: 0, + EditStrategies: []editor.StrategyConfig{ + { + Type: "remove_tool_result", + Params: map[string]interface{}{ + "keep_recent_n_tool_results": 0, + }, + }, + }, + PinEditingStrategiesAtMessage: pinMsgID.String(), + EditingTrigger: &EditingTrigger{ + TokenGte: &triggerThreshold, + }, + }) + assert.NoError(t, err) + assert.NotNil(t, out) + assert.Equal(t, pinMsgID.String(), out.EditAtMessageID) + assert.NotEqual(t, lastMsgID.String(), out.EditAtMessageID) + + repo.AssertExpectations(t) +} + +func TestSessionService_GetMessages_EditTriggerBranchCoverage(t *testing.T) { + ctx := context.Background() + sessionID := uuid.New() + now := time.Now() + + err := tokenizer.Init(zap.NewNop()) + assert.NoError(t, err) + + msgID1 := uuid.MustParse("00000000-0000-0000-0000-000000000101") + msgID2 := uuid.MustParse("00000000-0000-0000-0000-000000000102") + msgID3 := uuid.MustParse("00000000-0000-0000-0000-000000000103") + missingPinID := "00000000-0000-0000-0000-000000000999" + highThreshold := 100000 + lowThreshold := 1 + zeroThreshold := 0 + + makeToolResultMessages := func() []model.Message { + return []model.Message{ + { + ID: msgID1, + SessionID: sessionID, + Role: model.RoleAssistant, + CreatedAt: now.Add(1 * time.Minute), + Parts: []model.Part{ + { + Type: model.PartTypeToolResult, + Text: "payload-one", + }, + }, + }, + { + ID: msgID2, + SessionID: sessionID, + Role: model.RoleAssistant, + CreatedAt: now.Add(2 * time.Minute), + Parts: []model.Part{ + { + Type: model.PartTypeToolResult, + Text: "payload-two", + }, + }, + }, + { + ID: msgID3, + SessionID: sessionID, + Role: model.RoleAssistant, + CreatedAt: now.Add(3 * time.Minute), + Parts: []model.Part{ + { + Type: model.PartTypeToolResult, + Text: "payload-three", + }, + }, + }, + } + } + + makeUserMessages := func() []model.Message { + return []model.Message{ + { + ID: msgID1, + SessionID: sessionID, + Role: model.RoleUser, + CreatedAt: now.Add(1 * time.Minute), + Parts: []model.Part{ + model.NewTextPart("first"), + }, + }, + { + ID: msgID2, + SessionID: sessionID, + Role: model.RoleUser, + CreatedAt: now.Add(2 * time.Minute), + Parts: []model.Part{ + model.NewTextPart("second"), + }, + }, + { + ID: msgID3, + SessionID: sessionID, + Role: model.RoleUser, + CreatedAt: now.Add(3 * time.Minute), + Parts: []model.Part{ + model.NewTextPart("third"), + }, + }, + } + } + + editStrategies := []editor.StrategyConfig{ + { + Type: "remove_tool_result", + Params: map[string]interface{}{ + "keep_recent_n_tool_results": 0, + }, + }, + } + + tests := []struct { + name string + input GetMessagesInput + repoMode string // all|paged + repoMessages []model.Message + wantEditAt string + wantPartsText []string + wantItemsLen int + wantHasMore bool + wantNextCursorSet bool + }{ + { + name: "no edit strategies sets last message id", + input: GetMessagesInput{ + SessionID: sessionID, + Limit: 0, + }, + repoMode: "all", + repoMessages: makeUserMessages(), + wantEditAt: msgID3.String(), + wantItemsLen: 3, + wantPartsText: []string{"first", "second", "third"}, + }, + { + name: "no trigger configured applies strategies", + input: GetMessagesInput{ + SessionID: sessionID, + Limit: 0, + EditStrategies: editStrategies, + }, + repoMode: "all", + repoMessages: makeToolResultMessages(), + wantEditAt: msgID3.String(), + wantItemsLen: 3, + wantPartsText: []string{"Done", "Done", "Done"}, + }, + { + name: "trigger false without pin skips strategies and uses last message id", + input: GetMessagesInput{ + SessionID: sessionID, + Limit: 0, + EditStrategies: editStrategies, + EditingTrigger: &EditingTrigger{ + TokenGte: &highThreshold, + }, + }, + repoMode: "all", + repoMessages: makeToolResultMessages(), + wantEditAt: msgID3.String(), + wantItemsLen: 3, + wantPartsText: []string{"payload-one", "payload-two", "payload-three"}, + }, + { + name: "trigger false with pin preserves provided pin", + input: GetMessagesInput{ + SessionID: sessionID, + Limit: 0, + EditStrategies: editStrategies, + PinEditingStrategiesAtMessage: msgID1.String(), + EditingTrigger: &EditingTrigger{ + TokenGte: &highThreshold, + }, + }, + repoMode: "all", + repoMessages: makeToolResultMessages(), + wantEditAt: msgID1.String(), + wantItemsLen: 3, + wantPartsText: []string{"payload-one", "payload-two", "payload-three"}, + }, + { + name: "trigger true with pin found applies only up to pin message", + input: GetMessagesInput{ + SessionID: sessionID, + Limit: 0, + EditStrategies: editStrategies, + PinEditingStrategiesAtMessage: msgID1.String(), + EditingTrigger: &EditingTrigger{ + TokenGte: &lowThreshold, + }, + }, + repoMode: "all", + repoMessages: makeToolResultMessages(), + wantEditAt: msgID1.String(), + wantItemsLen: 3, + wantPartsText: []string{"Done", "payload-two", "payload-three"}, + }, + { + name: "trigger true with pin not found applies to all messages", + input: GetMessagesInput{ + SessionID: sessionID, + Limit: 0, + EditStrategies: editStrategies, + PinEditingStrategiesAtMessage: missingPinID, + EditingTrigger: &EditingTrigger{ + TokenGte: &lowThreshold, + }, + }, + repoMode: "all", + repoMessages: makeToolResultMessages(), + wantEditAt: msgID3.String(), + wantItemsLen: 3, + wantPartsText: []string{"Done", "Done", "Done"}, + }, + { + name: "non-positive token threshold behaves as trigger not configured", + input: GetMessagesInput{ + SessionID: sessionID, + Limit: 0, + EditStrategies: editStrategies, + EditingTrigger: &EditingTrigger{ + TokenGte: &zeroThreshold, + }, + }, + repoMode: "all", + repoMessages: makeToolResultMessages(), + wantEditAt: msgID3.String(), + wantItemsLen: 3, + wantPartsText: []string{"Done", "Done", "Done"}, + }, + { + name: "pagination window excludes pin and follows pin-not-found branch", + input: GetMessagesInput{ + SessionID: sessionID, + Limit: 2, + EditStrategies: editStrategies, + PinEditingStrategiesAtMessage: msgID3.String(), + EditingTrigger: &EditingTrigger{ + TokenGte: &lowThreshold, + }, + }, + repoMode: "paged", + repoMessages: makeToolResultMessages(), + wantEditAt: msgID2.String(), + wantItemsLen: 2, + wantPartsText: []string{"Done", "Done"}, + wantHasMore: true, + wantNextCursorSet: true, + }, + { + name: "empty messages with trigger false and pin keeps pin", + input: GetMessagesInput{ + SessionID: sessionID, + Limit: 0, + EditStrategies: editStrategies, + PinEditingStrategiesAtMessage: msgID1.String(), + EditingTrigger: &EditingTrigger{ + TokenGte: &highThreshold, + }, + }, + repoMode: "all", + repoMessages: []model.Message{}, + wantEditAt: msgID1.String(), + wantItemsLen: 0, + wantPartsText: []string{}, + }, + { + name: "empty messages with no strategies keeps edit_at_message_id empty", + input: GetMessagesInput{ + SessionID: sessionID, + Limit: 0, + }, + repoMode: "all", + repoMessages: []model.Message{}, + wantEditAt: "", + wantItemsLen: 0, + wantPartsText: []string{}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + repo := &MockSessionRepo{} + if tt.repoMode == "paged" { + repo.On("ListBySessionWithCursor", ctx, sessionID, time.Time{}, uuid.UUID{}, tt.input.Limit+1, tt.input.TimeDesc).Return(tt.repoMessages, nil) + } else { + repo.On("ListAllMessagesBySession", ctx, sessionID).Return(tt.repoMessages, nil) + } + + service := NewSessionService(repo, &MockAssetReferenceRepo{}, zap.NewNop(), nil, nil, &config.Config{}, nil) + out, err := service.GetMessages(ctx, tt.input) + assert.NoError(t, err) + assert.NotNil(t, out) + + assert.Equal(t, tt.wantEditAt, out.EditAtMessageID) + assert.Equal(t, tt.wantItemsLen, len(out.Items)) + assert.Equal(t, tt.wantHasMore, out.HasMore) + if tt.wantNextCursorSet { + assert.NotEmpty(t, out.NextCursor) + } + if !tt.wantNextCursorSet { + assert.Empty(t, out.NextCursor) + } + + gotPartsText := make([]string, 0, len(out.Items)) + for _, msg := range out.Items { + if len(msg.Parts) > 0 { + gotPartsText = append(gotPartsText, msg.Parts[0].Text) + } + } + assert.Equal(t, tt.wantPartsText, gotPartsText) + if len(out.Items) > 0 { + assert.Greater(t, out.ThisTimeTokens, 0) + } else { + assert.Equal(t, 0, out.ThisTimeTokens) + } + + repo.AssertExpectations(t) + }) + } +} + +func TestBuildEditingTriggerChecks_Branches(t *testing.T) { + ctx := context.Background() + + err := tokenizer.Init(zap.NewNop()) + assert.NoError(t, err) + + positive := 1 + zero := 0 + + msg := model.Message{ + ID: uuid.New(), + Role: model.RoleUser, + Parts: []model.Part{ + model.NewTextPart("hello world"), + }, + } + + t.Run("nil trigger returns no checks", func(t *testing.T) { + checks := buildEditingTriggerChecks(nil) + assert.Len(t, checks, 0) + }) + + t.Run("nil token_gte returns no checks", func(t *testing.T) { + checks := buildEditingTriggerChecks(&EditingTrigger{}) + assert.Len(t, checks, 0) + }) + + t.Run("non-positive token_gte returns no checks", func(t *testing.T) { + checks := buildEditingTriggerChecks(&EditingTrigger{TokenGte: &zero}) + assert.Len(t, checks, 0) + }) + + t.Run("positive token_gte adds a check and evaluates", func(t *testing.T) { + checks := buildEditingTriggerChecks(&EditingTrigger{TokenGte: &positive}) + assert.Len(t, checks, 1) + + eval := &editingTriggerEval{ + sessionID: uuid.New(), + messages: []model.Message{msg}, + } + + ok, err := checks[0](ctx, eval) + assert.NoError(t, err) + assert.True(t, ok) + assert.NotNil(t, eval.tokenCount) + }) +} + +func TestEditingTriggerEvalTokens_UsesCachedValue(t *testing.T) { + expected := 123 + eval := &editingTriggerEval{ + sessionID: uuid.New(), + messages: nil, + tokenCount: &expected, + } + + tokens, err := eval.Tokens(context.Background()) + assert.NoError(t, err) + assert.Equal(t, expected, tokens) +} + +func TestSameMessageOrderByID_Branches(t *testing.T) { + id1 := uuid.MustParse("00000000-0000-0000-0000-000000000201") + id2 := uuid.MustParse("00000000-0000-0000-0000-000000000202") + id3 := uuid.MustParse("00000000-0000-0000-0000-000000000203") + + t.Run("length mismatch returns false", func(t *testing.T) { + a := []model.Message{{ID: id1}} + b := []model.Message{{ID: id1}, {ID: id2}} + assert.False(t, sameMessageOrderByID(a, b)) + }) + + t.Run("id mismatch returns false", func(t *testing.T) { + a := []model.Message{{ID: id1}, {ID: id2}} + b := []model.Message{{ID: id1}, {ID: id3}} + assert.False(t, sameMessageOrderByID(a, b)) + }) + + t.Run("same order returns true", func(t *testing.T) { + a := []model.Message{{ID: id1}, {ID: id2}} + b := []model.Message{{ID: id1}, {ID: id2}} + assert.True(t, sameMessageOrderByID(a, b)) + }) +} From 9a336825a60510e660be22a6dffcee89e31b97d0 Mon Sep 17 00:00:00 2001 From: riturajFi Date: Sat, 21 Feb 2026 18:57:33 +0530 Subject: [PATCH 07/20] add codex plans to ignore --- .gitignore | 1 + 1 file changed, 1 insertion(+) diff --git a/.gitignore b/.gitignore index 6a7a10d44..5a2a6bb4b 100644 --- a/.gitignore +++ b/.gitignore @@ -1,6 +1,7 @@ plans/ .claude .cursorrules +plans/ .DS_Store .agents skills-lock.json From 817d5e3c624bb02407d9a29cd6e276939c160d30 Mon Sep 17 00:00:00 2001 From: riturajFi Date: Sat, 21 Feb 2026 19:30:41 +0530 Subject: [PATCH 08/20] refactor --- .../go/internal/modules/service/session.go | 96 ++++----------- .../internal/modules/service/session_test.go | 85 ------------- .../pkg/editingtrigger/check_registry.go | 14 +++ .../pkg/editingtrigger/editing_trigger.go | 86 +++++++++++++ .../editingtrigger/editing_trigger_test.go | 114 ++++++++++++++++++ .../pkg/editingtrigger/token_gte_check.go | 23 ++++ 6 files changed, 262 insertions(+), 156 deletions(-) create mode 100644 src/server/api/go/internal/pkg/editingtrigger/check_registry.go create mode 100644 src/server/api/go/internal/pkg/editingtrigger/editing_trigger.go create mode 100644 src/server/api/go/internal/pkg/editingtrigger/editing_trigger_test.go create mode 100644 src/server/api/go/internal/pkg/editingtrigger/token_gte_check.go diff --git a/src/server/api/go/internal/modules/service/session.go b/src/server/api/go/internal/modules/service/session.go index 73b5f3f38..66e17be3e 100644 --- a/src/server/api/go/internal/modules/service/session.go +++ b/src/server/api/go/internal/modules/service/session.go @@ -19,6 +19,7 @@ import ( mq "github.com/memodb-io/Acontext/internal/infra/queue" "github.com/memodb-io/Acontext/internal/modules/model" "github.com/memodb-io/Acontext/internal/modules/repo" + "github.com/memodb-io/Acontext/internal/pkg/editingtrigger" "github.com/memodb-io/Acontext/internal/pkg/editor" "github.com/memodb-io/Acontext/internal/pkg/paging" "github.com/memodb-io/Acontext/internal/pkg/tokenizer" @@ -454,67 +455,7 @@ type GetMessagesInput struct { UserKEK []byte `json:"-"` // optional: for envelope encryption (decrypting parts) } -// EditingTrigger defines trigger configuration for applying edit_strategies. -// v0 supports only token_gte. -type EditingTrigger struct { - // TokenGte triggers edit strategies when the token count is greater than or equal to this value. - TokenGte *int `json:"token_gte,omitempty"` -} - -type editingTriggerEval struct { - sessionID uuid.UUID - messages []model.Message - - tokenCount *int -} - -func (e *editingTriggerEval) Tokens(ctx context.Context) (int, error) { - if e.tokenCount != nil { - return *e.tokenCount, nil - } - - tokens, err := tokenizer.CountMessagePartsTokens(ctx, e.messages) - if err != nil { - return 0, fmt.Errorf("%w: failed to count tokens for editing_trigger session_id=%s: %v", ErrGetMessagesTokenCount, e.sessionID, err) - } - - e.tokenCount = &tokens - return tokens, nil -} - -type editingTriggerCheck func(ctx context.Context, eval *editingTriggerEval) (bool, error) - -func buildEditingTriggerChecks(trigger *EditingTrigger) []editingTriggerCheck { - checks := make([]editingTriggerCheck, 0, 1) - if trigger == nil { - return checks - } - - if trigger.TokenGte != nil && *trigger.TokenGte > 0 { - threshold := *trigger.TokenGte - checks = append(checks, func(ctx context.Context, eval *editingTriggerEval) (bool, error) { - tokens, err := eval.Tokens(ctx) - if err != nil { - return false, err - } - return tokens >= threshold, nil - }) - } - - return checks -} - -func sameMessageOrderByID(a, b []model.Message) bool { - if len(a) != len(b) { - return false - } - for i := range a { - if a[i].ID != b[i].ID { - return false - } - } - return true -} +type EditingTrigger = editingtrigger.Trigger type PublicURL struct { URL string `json:"url"` @@ -619,11 +560,11 @@ func (s *sessionService) GetMessages(ctx context.Context, in GetMessagesInput) ( } // Apply edit strategies if provided (before format conversion) - var triggerEval *editingTriggerEval + var triggerEval *editingtrigger.Eval strategiesApplied := false if len(in.EditStrategies) > 0 { applyEditStrategies := true - triggerChecks := buildEditingTriggerChecks(in.EditingTrigger) + triggerChecks := editingtrigger.BuildChecks(in.EditingTrigger) triggerEvaluated := len(triggerChecks) > 0 if triggerEvaluated { // Evaluate trigger on the same editable prefix used by pin_editing_strategies_at_message. @@ -643,7 +584,17 @@ func (s *sessionService) GetMessages(ctx context.Context, in GetMessagesInput) ( // OR semantics: apply when any trigger check passes. applyEditStrategies = false - eval := &editingTriggerEval{sessionID: in.SessionID, messages: triggerMessages} + eval := editingtrigger.NewEval( + in.SessionID, + triggerMessages, + func(ctx context.Context, messages []model.Message) (int, error) { + tokens, err := tokenizer.CountMessagePartsTokens(ctx, messages) + if err != nil { + return 0, fmt.Errorf("%w: failed to count tokens for editing_trigger session_id=%s: %v", ErrGetMessagesTokenCount, in.SessionID, err) + } + return tokens, nil + }, + ) triggerEval = eval for _, check := range triggerChecks { ok, err := check(ctx, eval) @@ -702,16 +653,19 @@ func (s *sessionService) GetMessages(ctx context.Context, in GetMessagesInput) ( } } - if triggerEval != nil && triggerEval.tokenCount != nil && !strategiesApplied && sameMessageOrderByID(triggerEval.messages, out.Items) { - out.ThisTimeTokens = *triggerEval.tokenCount - } else { - thisTimeTokens, err := tokenizer.CountMessagePartsTokens(ctx, out.Items) - if err != nil { - return nil, fmt.Errorf("%w: session_id=%s: %v", ErrGetMessagesTokenCount, in.SessionID, err) + if triggerEval != nil { + if cachedTokens, ok := triggerEval.CachedTokens(); ok && !strategiesApplied && editingtrigger.SameMessageOrderByID(triggerEval.Messages(), out.Items) { + out.ThisTimeTokens = cachedTokens + return out, nil } - out.ThisTimeTokens = thisTimeTokens } + thisTimeTokens, err := tokenizer.CountMessagePartsTokens(ctx, out.Items) + if err != nil { + return nil, fmt.Errorf("%w: session_id=%s: %v", ErrGetMessagesTokenCount, in.SessionID, err) + } + out.ThisTimeTokens = thisTimeTokens + return out, nil } diff --git a/src/server/api/go/internal/modules/service/session_test.go b/src/server/api/go/internal/modules/service/session_test.go index b8ad0b7d3..394f3f907 100644 --- a/src/server/api/go/internal/modules/service/session_test.go +++ b/src/server/api/go/internal/modules/service/session_test.go @@ -2161,88 +2161,3 @@ func TestSessionService_GetMessages_EditTriggerBranchCoverage(t *testing.T) { }) } } - -func TestBuildEditingTriggerChecks_Branches(t *testing.T) { - ctx := context.Background() - - err := tokenizer.Init(zap.NewNop()) - assert.NoError(t, err) - - positive := 1 - zero := 0 - - msg := model.Message{ - ID: uuid.New(), - Role: model.RoleUser, - Parts: []model.Part{ - model.NewTextPart("hello world"), - }, - } - - t.Run("nil trigger returns no checks", func(t *testing.T) { - checks := buildEditingTriggerChecks(nil) - assert.Len(t, checks, 0) - }) - - t.Run("nil token_gte returns no checks", func(t *testing.T) { - checks := buildEditingTriggerChecks(&EditingTrigger{}) - assert.Len(t, checks, 0) - }) - - t.Run("non-positive token_gte returns no checks", func(t *testing.T) { - checks := buildEditingTriggerChecks(&EditingTrigger{TokenGte: &zero}) - assert.Len(t, checks, 0) - }) - - t.Run("positive token_gte adds a check and evaluates", func(t *testing.T) { - checks := buildEditingTriggerChecks(&EditingTrigger{TokenGte: &positive}) - assert.Len(t, checks, 1) - - eval := &editingTriggerEval{ - sessionID: uuid.New(), - messages: []model.Message{msg}, - } - - ok, err := checks[0](ctx, eval) - assert.NoError(t, err) - assert.True(t, ok) - assert.NotNil(t, eval.tokenCount) - }) -} - -func TestEditingTriggerEvalTokens_UsesCachedValue(t *testing.T) { - expected := 123 - eval := &editingTriggerEval{ - sessionID: uuid.New(), - messages: nil, - tokenCount: &expected, - } - - tokens, err := eval.Tokens(context.Background()) - assert.NoError(t, err) - assert.Equal(t, expected, tokens) -} - -func TestSameMessageOrderByID_Branches(t *testing.T) { - id1 := uuid.MustParse("00000000-0000-0000-0000-000000000201") - id2 := uuid.MustParse("00000000-0000-0000-0000-000000000202") - id3 := uuid.MustParse("00000000-0000-0000-0000-000000000203") - - t.Run("length mismatch returns false", func(t *testing.T) { - a := []model.Message{{ID: id1}} - b := []model.Message{{ID: id1}, {ID: id2}} - assert.False(t, sameMessageOrderByID(a, b)) - }) - - t.Run("id mismatch returns false", func(t *testing.T) { - a := []model.Message{{ID: id1}, {ID: id2}} - b := []model.Message{{ID: id1}, {ID: id3}} - assert.False(t, sameMessageOrderByID(a, b)) - }) - - t.Run("same order returns true", func(t *testing.T) { - a := []model.Message{{ID: id1}, {ID: id2}} - b := []model.Message{{ID: id1}, {ID: id2}} - assert.True(t, sameMessageOrderByID(a, b)) - }) -} diff --git a/src/server/api/go/internal/pkg/editingtrigger/check_registry.go b/src/server/api/go/internal/pkg/editingtrigger/check_registry.go new file mode 100644 index 000000000..9c0f4c9a6 --- /dev/null +++ b/src/server/api/go/internal/pkg/editingtrigger/check_registry.go @@ -0,0 +1,14 @@ +package editingtrigger + +// checkBuilder returns a check and whether it should be included for a trigger. +type checkBuilder func(trigger *Trigger) (Check, bool) + +var registeredCheckBuilders []checkBuilder + +func registerCheckBuilder(builder checkBuilder) { + registeredCheckBuilders = append(registeredCheckBuilders, builder) +} + +func listCheckBuilders() []checkBuilder { + return registeredCheckBuilders +} diff --git a/src/server/api/go/internal/pkg/editingtrigger/editing_trigger.go b/src/server/api/go/internal/pkg/editingtrigger/editing_trigger.go new file mode 100644 index 000000000..2ad7d08d3 --- /dev/null +++ b/src/server/api/go/internal/pkg/editingtrigger/editing_trigger.go @@ -0,0 +1,86 @@ +package editingtrigger + +import ( + "context" + + "github.com/google/uuid" + "github.com/memodb-io/Acontext/internal/modules/model" +) + +// Trigger defines trigger configuration for applying edit strategies. +// v0 supports only token_gte. +type Trigger struct { + // TokenGte triggers edit strategies when token count is >= this value. + TokenGte *int `json:"token_gte,omitempty"` +} + +// TokenCounter computes token count for a message slice. +type TokenCounter func(ctx context.Context, messages []model.Message) (int, error) + +// Eval evaluates trigger checks and memoizes token count. +type Eval struct { + sessionID uuid.UUID + messages []model.Message + counter TokenCounter + + tokenCount *int +} + +func NewEval(sessionID uuid.UUID, messages []model.Message, counter TokenCounter) *Eval { + return &Eval{ + sessionID: sessionID, + messages: messages, + counter: counter, + } +} + +func (e *Eval) Tokens(ctx context.Context) (int, error) { + if e.tokenCount != nil { + return *e.tokenCount, nil + } + + tokens, err := e.counter(ctx, e.messages) + if err != nil { + return 0, err + } + + e.tokenCount = &tokens + return tokens, nil +} + +func (e *Eval) Messages() []model.Message { + return e.messages +} + +func (e *Eval) CachedTokens() (int, bool) { + if e.tokenCount == nil { + return 0, false + } + return *e.tokenCount, true +} + +type Check func(ctx context.Context, eval *Eval) (bool, error) + +func BuildChecks(trigger *Trigger) []Check { + builders := listCheckBuilders() + checks := make([]Check, 0, len(builders)) + for _, builder := range builders { + check, ok := builder(trigger) + if ok { + checks = append(checks, check) + } + } + return checks +} + +func SameMessageOrderByID(a, b []model.Message) bool { + if len(a) != len(b) { + return false + } + for i := range a { + if a[i].ID != b[i].ID { + return false + } + } + return true +} diff --git a/src/server/api/go/internal/pkg/editingtrigger/editing_trigger_test.go b/src/server/api/go/internal/pkg/editingtrigger/editing_trigger_test.go new file mode 100644 index 000000000..810d31760 --- /dev/null +++ b/src/server/api/go/internal/pkg/editingtrigger/editing_trigger_test.go @@ -0,0 +1,114 @@ +package editingtrigger + +import ( + "context" + "errors" + "testing" + + "github.com/google/uuid" + "github.com/memodb-io/Acontext/internal/modules/model" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestBuildChecks_Branches(t *testing.T) { + ctx := context.Background() + + positive := 1 + zero := 0 + + msg := model.Message{ + ID: uuid.New(), + Role: model.RoleUser, + Parts: []model.Part{ + model.NewTextPart("hello world"), + }, + } + + t.Run("nil trigger returns no checks", func(t *testing.T) { + checks := BuildChecks(nil) + assert.Len(t, checks, 0) + }) + + t.Run("nil token_gte returns no checks", func(t *testing.T) { + checks := BuildChecks(&Trigger{}) + assert.Len(t, checks, 0) + }) + + t.Run("non-positive token_gte returns no checks", func(t *testing.T) { + checks := BuildChecks(&Trigger{TokenGte: &zero}) + assert.Len(t, checks, 0) + }) + + t.Run("positive token_gte adds a check and evaluates", func(t *testing.T) { + calls := 0 + checks := BuildChecks(&Trigger{TokenGte: &positive}) + assert.Len(t, checks, 1) + + eval := NewEval(uuid.New(), []model.Message{msg}, func(context.Context, []model.Message) (int, error) { + calls++ + return 5, nil + }) + + ok, err := checks[0](ctx, eval) + assert.NoError(t, err) + assert.True(t, ok) + assert.Equal(t, 1, calls) + + tokens, cached := eval.CachedTokens() + require.True(t, cached) + assert.Equal(t, 5, tokens) + }) +} + +func TestEvalTokens_UsesCachedValue(t *testing.T) { + calls := 0 + eval := NewEval(uuid.New(), nil, func(context.Context, []model.Message) (int, error) { + calls++ + return 123, nil + }) + + first, err := eval.Tokens(context.Background()) + require.NoError(t, err) + second, err := eval.Tokens(context.Background()) + require.NoError(t, err) + + assert.Equal(t, 123, first) + assert.Equal(t, 123, second) + assert.Equal(t, 1, calls) +} + +func TestEvalTokens_PropagatesCounterError(t *testing.T) { + expectedErr := errors.New("counter failed") + eval := NewEval(uuid.New(), nil, func(context.Context, []model.Message) (int, error) { + return 0, expectedErr + }) + + _, err := eval.Tokens(context.Background()) + require.Error(t, err) + assert.ErrorIs(t, err, expectedErr) +} + +func TestSameMessageOrderByID_Branches(t *testing.T) { + id1 := uuid.MustParse("00000000-0000-0000-0000-000000000201") + id2 := uuid.MustParse("00000000-0000-0000-0000-000000000202") + id3 := uuid.MustParse("00000000-0000-0000-0000-000000000203") + + t.Run("length mismatch returns false", func(t *testing.T) { + a := []model.Message{{ID: id1}} + b := []model.Message{{ID: id1}, {ID: id2}} + assert.False(t, SameMessageOrderByID(a, b)) + }) + + t.Run("id mismatch returns false", func(t *testing.T) { + a := []model.Message{{ID: id1}, {ID: id2}} + b := []model.Message{{ID: id1}, {ID: id3}} + assert.False(t, SameMessageOrderByID(a, b)) + }) + + t.Run("same order returns true", func(t *testing.T) { + a := []model.Message{{ID: id1}, {ID: id2}} + b := []model.Message{{ID: id1}, {ID: id2}} + assert.True(t, SameMessageOrderByID(a, b)) + }) +} diff --git a/src/server/api/go/internal/pkg/editingtrigger/token_gte_check.go b/src/server/api/go/internal/pkg/editingtrigger/token_gte_check.go new file mode 100644 index 000000000..25260ad63 --- /dev/null +++ b/src/server/api/go/internal/pkg/editingtrigger/token_gte_check.go @@ -0,0 +1,23 @@ +package editingtrigger + +import "context" + +func init() { + registerCheckBuilder(buildTokenGteCheck) +} + +func buildTokenGteCheck(trigger *Trigger) (Check, bool) { + if trigger == nil || trigger.TokenGte == nil || *trigger.TokenGte <= 0 { + return nil, false + } + + threshold := *trigger.TokenGte + check := func(ctx context.Context, eval *Eval) (bool, error) { + tokens, err := eval.Tokens(ctx) + if err != nil { + return false, err + } + return tokens >= threshold, nil + } + return check, true +} From af109d7d646775c8194b6fdb52e21b4e8183053a Mon Sep 17 00:00:00 2001 From: riturajFi Date: Thu, 5 Mar 2026 21:22:07 +0530 Subject: [PATCH 09/20] fix 1 --- .env.local-api | 30 ------------------------------ .env.local-api.example | 30 ++++++++++++++++++++++++++++++ .gitignore | 2 ++ src/server/.env.local-api | 30 ------------------------------ src/server/.env.local-api.example | 30 ++++++++++++++++++++++++++++++ 5 files changed, 62 insertions(+), 60 deletions(-) delete mode 100644 .env.local-api create mode 100644 .env.local-api.example delete mode 100644 src/server/.env.local-api create mode 100644 src/server/.env.local-api.example diff --git a/.env.local-api b/.env.local-api deleted file mode 100644 index 05ea7dab6..000000000 --- a/.env.local-api +++ /dev/null @@ -1,30 +0,0 @@ -APP_ENV=debug -API_EXPORT_PORT=8029 -ROOT_API_BEARER_TOKEN=local-dev-root - -DATABASE_HOST=127.0.0.1 -DATABASE_EXPORT_PORT=15432 -DATABASE_USER=acontext -DATABASE_PASSWORD=helloworld -DATABASE_NAME=acontext - -REDIS_HOST=127.0.0.1 -REDIS_EXPORT_PORT=16379 -REDIS_PASSWORD=helloworld - -RABBITMQ_HOST=127.0.0.1 -RABBITMQ_EXPORT_PORT=15672 -RABBITMQ_USER=acontext -RABBITMQ_PASSWORD=helloworld -RABBITMQ_VHOST=/ -RABBITMQ_VHOST_ENCODED=%2F - -S3_ENDPOINT=http://127.0.0.1:19000 -S3_INTERNAL_ENDPOINT=http://127.0.0.1:19000 -S3_REGION=auto -S3_ACCESS_KEY=acontext -S3_SECRET_KEY=helloworld -S3_BUCKET=acontext-assets - -CORE_BASE_URL=http://127.0.0.1:8019 -OTEL_EXPORTER_OTLP_ENDPOINT= diff --git a/.env.local-api.example b/.env.local-api.example new file mode 100644 index 000000000..6574e17ee --- /dev/null +++ b/.env.local-api.example @@ -0,0 +1,30 @@ +APP_ENV=debug +API_EXPORT_PORT=8029 +ROOT_API_BEARER_TOKEN=replace-with-local-dev-token + +DATABASE_HOST=127.0.0.1 +DATABASE_EXPORT_PORT=15432 +DATABASE_USER=replace-with-db-user +DATABASE_PASSWORD=replace-with-db-password +DATABASE_NAME=replace-with-db-name + +REDIS_HOST=127.0.0.1 +REDIS_EXPORT_PORT=16379 +REDIS_PASSWORD=replace-with-redis-password + +RABBITMQ_HOST=127.0.0.1 +RABBITMQ_EXPORT_PORT=15672 +RABBITMQ_USER=replace-with-rabbitmq-user +RABBITMQ_PASSWORD=replace-with-rabbitmq-password +RABBITMQ_VHOST=/ +RABBITMQ_VHOST_ENCODED=%2F + +S3_ENDPOINT=http://127.0.0.1:19000 +S3_INTERNAL_ENDPOINT=http://127.0.0.1:19000 +S3_REGION=auto +S3_ACCESS_KEY=replace-with-s3-access-key +S3_SECRET_KEY=replace-with-s3-secret-key +S3_BUCKET=replace-with-s3-bucket + +CORE_BASE_URL=http://127.0.0.1:8019 +OTEL_EXPORTER_OTLP_ENDPOINT= diff --git a/.gitignore b/.gitignore index 5a2a6bb4b..ae49c2857 100644 --- a/.gitignore +++ b/.gitignore @@ -2,6 +2,8 @@ plans/ .claude .cursorrules plans/ +.env.local-api +src/server/.env.local-api .DS_Store .agents skills-lock.json diff --git a/src/server/.env.local-api b/src/server/.env.local-api deleted file mode 100644 index 05ea7dab6..000000000 --- a/src/server/.env.local-api +++ /dev/null @@ -1,30 +0,0 @@ -APP_ENV=debug -API_EXPORT_PORT=8029 -ROOT_API_BEARER_TOKEN=local-dev-root - -DATABASE_HOST=127.0.0.1 -DATABASE_EXPORT_PORT=15432 -DATABASE_USER=acontext -DATABASE_PASSWORD=helloworld -DATABASE_NAME=acontext - -REDIS_HOST=127.0.0.1 -REDIS_EXPORT_PORT=16379 -REDIS_PASSWORD=helloworld - -RABBITMQ_HOST=127.0.0.1 -RABBITMQ_EXPORT_PORT=15672 -RABBITMQ_USER=acontext -RABBITMQ_PASSWORD=helloworld -RABBITMQ_VHOST=/ -RABBITMQ_VHOST_ENCODED=%2F - -S3_ENDPOINT=http://127.0.0.1:19000 -S3_INTERNAL_ENDPOINT=http://127.0.0.1:19000 -S3_REGION=auto -S3_ACCESS_KEY=acontext -S3_SECRET_KEY=helloworld -S3_BUCKET=acontext-assets - -CORE_BASE_URL=http://127.0.0.1:8019 -OTEL_EXPORTER_OTLP_ENDPOINT= diff --git a/src/server/.env.local-api.example b/src/server/.env.local-api.example new file mode 100644 index 000000000..6574e17ee --- /dev/null +++ b/src/server/.env.local-api.example @@ -0,0 +1,30 @@ +APP_ENV=debug +API_EXPORT_PORT=8029 +ROOT_API_BEARER_TOKEN=replace-with-local-dev-token + +DATABASE_HOST=127.0.0.1 +DATABASE_EXPORT_PORT=15432 +DATABASE_USER=replace-with-db-user +DATABASE_PASSWORD=replace-with-db-password +DATABASE_NAME=replace-with-db-name + +REDIS_HOST=127.0.0.1 +REDIS_EXPORT_PORT=16379 +REDIS_PASSWORD=replace-with-redis-password + +RABBITMQ_HOST=127.0.0.1 +RABBITMQ_EXPORT_PORT=15672 +RABBITMQ_USER=replace-with-rabbitmq-user +RABBITMQ_PASSWORD=replace-with-rabbitmq-password +RABBITMQ_VHOST=/ +RABBITMQ_VHOST_ENCODED=%2F + +S3_ENDPOINT=http://127.0.0.1:19000 +S3_INTERNAL_ENDPOINT=http://127.0.0.1:19000 +S3_REGION=auto +S3_ACCESS_KEY=replace-with-s3-access-key +S3_SECRET_KEY=replace-with-s3-secret-key +S3_BUCKET=replace-with-s3-bucket + +CORE_BASE_URL=http://127.0.0.1:8019 +OTEL_EXPORTER_OTLP_ENDPOINT= From 68087410c6fbdbf851c0920ba57debbb77710a61 Mon Sep 17 00:00:00 2001 From: riturajFi Date: Thu, 5 Mar 2026 21:27:25 +0530 Subject: [PATCH 10/20] simplify edits --- .../pkg/editingtrigger/check_registry.go | 14 ----------- .../pkg/editingtrigger/editing_trigger.go | 22 ++++++++++++------ .../pkg/editingtrigger/token_gte_check.go | 23 ------------------- 3 files changed, 15 insertions(+), 44 deletions(-) delete mode 100644 src/server/api/go/internal/pkg/editingtrigger/check_registry.go delete mode 100644 src/server/api/go/internal/pkg/editingtrigger/token_gte_check.go diff --git a/src/server/api/go/internal/pkg/editingtrigger/check_registry.go b/src/server/api/go/internal/pkg/editingtrigger/check_registry.go deleted file mode 100644 index 9c0f4c9a6..000000000 --- a/src/server/api/go/internal/pkg/editingtrigger/check_registry.go +++ /dev/null @@ -1,14 +0,0 @@ -package editingtrigger - -// checkBuilder returns a check and whether it should be included for a trigger. -type checkBuilder func(trigger *Trigger) (Check, bool) - -var registeredCheckBuilders []checkBuilder - -func registerCheckBuilder(builder checkBuilder) { - registeredCheckBuilders = append(registeredCheckBuilders, builder) -} - -func listCheckBuilders() []checkBuilder { - return registeredCheckBuilders -} diff --git a/src/server/api/go/internal/pkg/editingtrigger/editing_trigger.go b/src/server/api/go/internal/pkg/editingtrigger/editing_trigger.go index 2ad7d08d3..980a198e6 100644 --- a/src/server/api/go/internal/pkg/editingtrigger/editing_trigger.go +++ b/src/server/api/go/internal/pkg/editingtrigger/editing_trigger.go @@ -62,14 +62,22 @@ func (e *Eval) CachedTokens() (int, bool) { type Check func(ctx context.Context, eval *Eval) (bool, error) func BuildChecks(trigger *Trigger) []Check { - builders := listCheckBuilders() - checks := make([]Check, 0, len(builders)) - for _, builder := range builders { - check, ok := builder(trigger) - if ok { - checks = append(checks, check) - } + if trigger == nil { + return nil + } + + checks := make([]Check, 0, 1) + if trigger.TokenGte != nil && *trigger.TokenGte > 0 { + threshold := *trigger.TokenGte + checks = append(checks, func(ctx context.Context, eval *Eval) (bool, error) { + tokens, err := eval.Tokens(ctx) + if err != nil { + return false, err + } + return tokens >= threshold, nil + }) } + return checks } diff --git a/src/server/api/go/internal/pkg/editingtrigger/token_gte_check.go b/src/server/api/go/internal/pkg/editingtrigger/token_gte_check.go deleted file mode 100644 index 25260ad63..000000000 --- a/src/server/api/go/internal/pkg/editingtrigger/token_gte_check.go +++ /dev/null @@ -1,23 +0,0 @@ -package editingtrigger - -import "context" - -func init() { - registerCheckBuilder(buildTokenGteCheck) -} - -func buildTokenGteCheck(trigger *Trigger) (Check, bool) { - if trigger == nil || trigger.TokenGte == nil || *trigger.TokenGte <= 0 { - return nil, false - } - - threshold := *trigger.TokenGte - check := func(ctx context.Context, eval *Eval) (bool, error) { - tokens, err := eval.Tokens(ctx) - if err != nil { - return false, err - } - return tokens >= threshold, nil - } - return check, true -} From f4c0978e732c4e965bc5fdc2408ae6efde91e78d Mon Sep 17 00:00:00 2001 From: riturajFi Date: Thu, 5 Mar 2026 21:48:34 +0530 Subject: [PATCH 11/20] minor refactor --- .codex | 0 .../internal/modules/handler/session_test.go | 6 +++- .../go/internal/modules/service/session.go | 14 +++++++- .../internal/modules/service/session_test.go | 32 ++++++++++++++++--- .../pkg/editingtrigger/editing_trigger.go | 12 ------- .../editingtrigger/editing_trigger_test.go | 24 -------------- 6 files changed, 46 insertions(+), 42 deletions(-) create mode 100644 .codex diff --git a/.codex b/.codex new file mode 100644 index 000000000..e69de29bb diff --git a/src/server/api/go/internal/modules/handler/session_test.go b/src/server/api/go/internal/modules/handler/session_test.go index afcfb331e..598678fc3 100644 --- a/src/server/api/go/internal/modules/handler/session_test.go +++ b/src/server/api/go/internal/modules/handler/session_test.go @@ -4357,12 +4357,16 @@ func TestSessionHandler_GetMessages_UsesServiceThisTimeTokens(t *testing.T) { func TestSessionHandler_GetMessages_RejectsEmptyEditingTrigger(t *testing.T) { gin.SetMode(gin.TestMode) + projectID := uuid.New() sessionID := uuid.New() mockService := &MockSessionService{} handler := NewSessionHandler(mockService, &MockUserService{}, getMockSessionCoreClient()) router := setupSessionRouter() - router.GET("/session/:session_id/messages", handler.GetMessages) + router.GET("/session/:session_id/messages", func(c *gin.Context) { + c.Set("project", &model.Project{ID: projectID}) + handler.GetMessages(c) + }) editStrategies := `[{"type":"token_limit","params":{"limit_tokens":100}}]` reqURL := "/session/" + sessionID.String() + "/messages?limit=20&edit_strategies=" + diff --git a/src/server/api/go/internal/modules/service/session.go b/src/server/api/go/internal/modules/service/session.go index 66e17be3e..457bab85f 100644 --- a/src/server/api/go/internal/modules/service/session.go +++ b/src/server/api/go/internal/modules/service/session.go @@ -654,7 +654,7 @@ func (s *sessionService) GetMessages(ctx context.Context, in GetMessagesInput) ( } if triggerEval != nil { - if cachedTokens, ok := triggerEval.CachedTokens(); ok && !strategiesApplied && editingtrigger.SameMessageOrderByID(triggerEval.Messages(), out.Items) { + if cachedTokens, ok := triggerEval.CachedTokens(); ok && !strategiesApplied && sameMessageOrderByID(triggerEval.Messages(), out.Items) { out.ThisTimeTokens = cachedTokens return out, nil } @@ -677,6 +677,18 @@ func (s *sessionService) DownloadAsset(ctx context.Context, s3Key string, userKE return s.s3.DownloadFile(ctx, s3Key, userKEK) } +func sameMessageOrderByID(a, b []model.Message) bool { + if len(a) != len(b) { + return false + } + for i := range a { + if a[i].ID != b[i].ID { + return false + } + } + return true +} + // cachePartsInRedis stores message parts in Redis with a fixed TTL. // When userKEK is provided, the serialized JSON is encrypted before caching. // Format: prefix_byte | payload diff --git a/src/server/api/go/internal/modules/service/session_test.go b/src/server/api/go/internal/modules/service/session_test.go index 394f3f907..f083ba2c4 100644 --- a/src/server/api/go/internal/modules/service/session_test.go +++ b/src/server/api/go/internal/modules/service/session_test.go @@ -1720,7 +1720,7 @@ func TestSessionService_GetMessages_ThisTimeTokensMatchesEditedOutput(t *testing assert.NoError(t, err) assert.Greater(t, preEditTokens, 0) - service := NewSessionService(repo, &MockAssetReferenceRepo{}, zap.NewNop(), nil, nil, &config.Config{}, nil) + service := NewSessionService(repo, nil, &MockAssetReferenceRepo{}, nil, zap.NewNop(), nil, nil, &config.Config{}, nil, nil) out, err := service.GetMessages(ctx, GetMessagesInput{ SessionID: sessionID, Limit: 0, @@ -1772,7 +1772,7 @@ func TestSessionService_GetMessages_TriggerTokenErrorsAreWrapped(t *testing.T) { }, }, nil) - service := NewSessionService(repo, &MockAssetReferenceRepo{}, zap.NewNop(), nil, nil, &config.Config{}, nil) + service := NewSessionService(repo, nil, &MockAssetReferenceRepo{}, nil, zap.NewNop(), nil, nil, &config.Config{}, nil, nil) out, err := service.GetMessages(ctx, GetMessagesInput{ SessionID: sessionID, Limit: 0, @@ -1828,7 +1828,7 @@ func TestSessionService_GetMessages_TriggerFalseKeepsProvidedPin(t *testing.T) { repo := &MockSessionRepo{} repo.On("ListAllMessagesBySession", ctx, sessionID).Return(repoMessages, nil) - service := NewSessionService(repo, &MockAssetReferenceRepo{}, zap.NewNop(), nil, nil, &config.Config{}, nil) + service := NewSessionService(repo, nil, &MockAssetReferenceRepo{}, nil, zap.NewNop(), nil, nil, &config.Config{}, nil, nil) out, err := service.GetMessages(ctx, GetMessagesInput{ SessionID: sessionID, Limit: 0, @@ -2129,7 +2129,7 @@ func TestSessionService_GetMessages_EditTriggerBranchCoverage(t *testing.T) { repo.On("ListAllMessagesBySession", ctx, sessionID).Return(tt.repoMessages, nil) } - service := NewSessionService(repo, &MockAssetReferenceRepo{}, zap.NewNop(), nil, nil, &config.Config{}, nil) + service := NewSessionService(repo, nil, &MockAssetReferenceRepo{}, nil, zap.NewNop(), nil, nil, &config.Config{}, nil, nil) out, err := service.GetMessages(ctx, tt.input) assert.NoError(t, err) assert.NotNil(t, out) @@ -2161,3 +2161,27 @@ func TestSessionService_GetMessages_EditTriggerBranchCoverage(t *testing.T) { }) } } + +func TestSameMessageOrderByID_Branches(t *testing.T) { + id1 := uuid.MustParse("00000000-0000-0000-0000-000000000201") + id2 := uuid.MustParse("00000000-0000-0000-0000-000000000202") + id3 := uuid.MustParse("00000000-0000-0000-0000-000000000203") + + t.Run("length mismatch returns false", func(t *testing.T) { + a := []model.Message{{ID: id1}} + b := []model.Message{{ID: id1}, {ID: id2}} + assert.False(t, sameMessageOrderByID(a, b)) + }) + + t.Run("id mismatch returns false", func(t *testing.T) { + a := []model.Message{{ID: id1}, {ID: id2}} + b := []model.Message{{ID: id1}, {ID: id3}} + assert.False(t, sameMessageOrderByID(a, b)) + }) + + t.Run("same order returns true", func(t *testing.T) { + a := []model.Message{{ID: id1}, {ID: id2}} + b := []model.Message{{ID: id1}, {ID: id2}} + assert.True(t, sameMessageOrderByID(a, b)) + }) +} diff --git a/src/server/api/go/internal/pkg/editingtrigger/editing_trigger.go b/src/server/api/go/internal/pkg/editingtrigger/editing_trigger.go index 980a198e6..009cd9d22 100644 --- a/src/server/api/go/internal/pkg/editingtrigger/editing_trigger.go +++ b/src/server/api/go/internal/pkg/editingtrigger/editing_trigger.go @@ -80,15 +80,3 @@ func BuildChecks(trigger *Trigger) []Check { return checks } - -func SameMessageOrderByID(a, b []model.Message) bool { - if len(a) != len(b) { - return false - } - for i := range a { - if a[i].ID != b[i].ID { - return false - } - } - return true -} diff --git a/src/server/api/go/internal/pkg/editingtrigger/editing_trigger_test.go b/src/server/api/go/internal/pkg/editingtrigger/editing_trigger_test.go index 810d31760..160009d2f 100644 --- a/src/server/api/go/internal/pkg/editingtrigger/editing_trigger_test.go +++ b/src/server/api/go/internal/pkg/editingtrigger/editing_trigger_test.go @@ -88,27 +88,3 @@ func TestEvalTokens_PropagatesCounterError(t *testing.T) { require.Error(t, err) assert.ErrorIs(t, err, expectedErr) } - -func TestSameMessageOrderByID_Branches(t *testing.T) { - id1 := uuid.MustParse("00000000-0000-0000-0000-000000000201") - id2 := uuid.MustParse("00000000-0000-0000-0000-000000000202") - id3 := uuid.MustParse("00000000-0000-0000-0000-000000000203") - - t.Run("length mismatch returns false", func(t *testing.T) { - a := []model.Message{{ID: id1}} - b := []model.Message{{ID: id1}, {ID: id2}} - assert.False(t, SameMessageOrderByID(a, b)) - }) - - t.Run("id mismatch returns false", func(t *testing.T) { - a := []model.Message{{ID: id1}, {ID: id2}} - b := []model.Message{{ID: id1}, {ID: id3}} - assert.False(t, SameMessageOrderByID(a, b)) - }) - - t.Run("same order returns true", func(t *testing.T) { - a := []model.Message{{ID: id1}, {ID: id2}} - b := []model.Message{{ID: id1}, {ID: id2}} - assert.True(t, SameMessageOrderByID(a, b)) - }) -} From 14521f004e24711fa1a180d5e5f82022e9d739c1 Mon Sep 17 00:00:00 2001 From: riturajFi Date: Fri, 6 Mar 2026 12:13:27 +0530 Subject: [PATCH 12/20] issue 2 --- .../go/internal/modules/handler/session.go | 50 +++------- .../internal/modules/handler/session_test.go | 96 +++++++++++++++++++ .../pkg/editingtrigger/editing_trigger.go | 59 ++++++++++++ .../editingtrigger/editing_trigger_test.go | 51 ++++++++++ 4 files changed, 218 insertions(+), 38 deletions(-) diff --git a/src/server/api/go/internal/modules/handler/session.go b/src/server/api/go/internal/modules/handler/session.go index 46812c605..9f71c35b5 100644 --- a/src/server/api/go/internal/modules/handler/session.go +++ b/src/server/api/go/internal/modules/handler/session.go @@ -19,6 +19,7 @@ import ( "github.com/memodb-io/Acontext/internal/modules/serializer" "github.com/memodb-io/Acontext/internal/modules/service" "github.com/memodb-io/Acontext/internal/pkg/converter" + "github.com/memodb-io/Acontext/internal/pkg/editingtrigger" "github.com/memodb-io/Acontext/internal/pkg/editor" "github.com/memodb-io/Acontext/internal/pkg/normalizer" "github.com/memodb-io/Acontext/internal/pkg/tokenizer" @@ -561,51 +562,24 @@ func (h *SessionHandler) GetMessages(c *gin.Context) { c.JSON(http.StatusBadRequest, serializer.ParamErr("editing_trigger requires edit_strategies", errors.New("missing edit_strategies"))) return } - - var raw map[string]interface{} - if err := sonic.Unmarshal([]byte(req.EditingTrigger), &raw); err != nil { - c.JSON(http.StatusBadRequest, serializer.ParamErr("invalid editing_trigger JSON", err)) - return - } - if len(raw) == 0 { - c.JSON(http.StatusBadRequest, serializer.ParamErr("invalid editing_trigger", errors.New("at least one supported trigger is required"))) - return - } - allowedTriggerKeys := map[string]struct{}{ - "token_gte": {}, - } - for k := range raw { - if _, ok := allowedTriggerKeys[k]; !ok { - c.JSON(http.StatusBadRequest, serializer.ParamErr("invalid editing_trigger", fmt.Errorf("unsupported trigger: %s", k))) + allowedTriggerKeys := map[string]struct{}{if err := sonic.Unmarshal([]byte(req.EditingTrigger), &raw); err != nil { c.JSON(http.StatusBadRequest, serializer.ParamErr("invalid editing_trigger JSON", err)) return } + var trig service.EditingTrigger + if err := json.Unmarshal([]byte(req.EditingTrigger), &trig); err != nil { + var unsupportedErr editingtrigger.UnsupportedTriggerError + if errors.As(err, &unsupportedErr) { + c.JSON(http.StatusBadRequest, serializer.ParamErr("invalid editing_trigger", err)) return } - } - - var trig service.EditingTrigger - if err := sonic.Unmarshal([]byte(req.EditingTrigger), &trig); err != nil { c.JSON(http.StatusBadRequest, serializer.ParamErr("invalid editing_trigger JSON", err)) return } - - triggerValidators := map[string]func(service.EditingTrigger) error{ - "token_gte": func(t service.EditingTrigger) error { - if t.TokenGte == nil || *t.TokenGte <= 0 { - return errors.New("token_gte must be > 0") - } - return nil - }, - } - for k := range raw { - validate, ok := triggerValidators[k] - if !ok { - // Should be unreachable due to allowedTriggerKeys check above. - c.JSON(http.StatusBadRequest, serializer.ParamErr("invalid editing_trigger", fmt.Errorf("unsupported trigger: %s", k))) - return - } - if err := validate(trig); err != nil { - c.JSON(http.StatusBadRequest, serializer.ParamErr("invalid editing_trigger."+k, err)) + if err := trig.Validate(); err != nil { + if errors.Is(err, editingtrigger.ErrTokenGteMustBeGreater) { + c.JSON(http.StatusBadRequest, serializer.ParamErr("invalid editing_trigger.token_gte", err)) return } + c.JSON(http.StatusBadRequest, serializer.ParamErr("invalid editing_trigger", err)) + return } editingTrigger = &trig } diff --git a/src/server/api/go/internal/modules/handler/session_test.go b/src/server/api/go/internal/modules/handler/session_test.go index 598678fc3..7c5ba09c0 100644 --- a/src/server/api/go/internal/modules/handler/session_test.go +++ b/src/server/api/go/internal/modules/handler/session_test.go @@ -4378,3 +4378,99 @@ func TestSessionHandler_GetMessages_RejectsEmptyEditingTrigger(t *testing.T) { assert.Equal(t, http.StatusBadRequest, w.Code) mockService.AssertNotCalled(t, "GetMessages") } + +func TestSessionHandler_GetMessages_RejectsUnknownEditingTrigger(t *testing.T) { + gin.SetMode(gin.TestMode) + + sessionID := uuid.New() + mockService := &MockSessionService{} + + handler := NewSessionHandler(mockService, &MockUserService{}, getMockSessionCoreClient()) + router := setupSessionRouter() + router.GET("/session/:session_id/messages", handler.GetMessages) + + editStrategies := `[{"type":"token_limit","params":{"limit_tokens":100}}]` + reqURL := "/session/" + sessionID.String() + "/messages?limit=20&edit_strategies=" + + url.QueryEscape(editStrategies) + "&editing_trigger=" + url.QueryEscape(`{"unknown":1}`) + req := httptest.NewRequest("GET", reqURL, nil) + w := httptest.NewRecorder() + router.ServeHTTP(w, req) + + assert.Equal(t, http.StatusBadRequest, w.Code) + mockService.AssertNotCalled(t, "GetMessages") +} + +func TestSessionHandler_GetMessages_RejectsInvalidTokenGte(t *testing.T) { + gin.SetMode(gin.TestMode) + + sessionID := uuid.New() + mockService := &MockSessionService{} + + handler := NewSessionHandler(mockService, &MockUserService{}, getMockSessionCoreClient()) + router := setupSessionRouter() + router.GET("/session/:session_id/messages", handler.GetMessages) + + editStrategies := `[{"type":"token_limit","params":{"limit_tokens":100}}]` + reqURL := "/session/" + sessionID.String() + "/messages?limit=20&edit_strategies=" + + url.QueryEscape(editStrategies) + "&editing_trigger=" + url.QueryEscape(`{"token_gte":0}`) + req := httptest.NewRequest("GET", reqURL, nil) + w := httptest.NewRecorder() + router.ServeHTTP(w, req) + + assert.Equal(t, http.StatusBadRequest, w.Code) + mockService.AssertNotCalled(t, "GetMessages") +} + +func TestSessionHandler_GetMessages_RejectsTriggerWithoutStrategies(t *testing.T) { + gin.SetMode(gin.TestMode) + + sessionID := uuid.New() + mockService := &MockSessionService{} + + handler := NewSessionHandler(mockService, &MockUserService{}, getMockSessionCoreClient()) + router := setupSessionRouter() + router.GET("/session/:session_id/messages", handler.GetMessages) + + reqURL := "/session/" + sessionID.String() + "/messages?limit=20&editing_trigger=" + url.QueryEscape(`{"token_gte":1}`) + req := httptest.NewRequest("GET", reqURL, nil) + w := httptest.NewRecorder() + router.ServeHTTP(w, req) + + assert.Equal(t, http.StatusBadRequest, w.Code) + mockService.AssertNotCalled(t, "GetMessages") +} + +func TestSessionHandler_GetMessages_AcceptsValidEditingTrigger(t *testing.T) { + gin.SetMode(gin.TestMode) + + sessionID := uuid.New() + mockService := &MockSessionService{} + mockService.On( + "GetMessages", + mock.Anything, + mock.MatchedBy(func(in service.GetMessagesInput) bool { + if in.EditingTrigger == nil || in.EditingTrigger.TokenGte == nil { + return false + } + return *in.EditingTrigger.TokenGte == 30000 + }), + ).Return(&service.GetMessagesOutput{ + Items: []model.Message{}, + HasMore: false, + ThisTimeTokens: 0, + }, nil) + + handler := NewSessionHandler(mockService, &MockUserService{}, getMockSessionCoreClient()) + router := setupSessionRouter() + router.GET("/session/:session_id/messages", handler.GetMessages) + + editStrategies := `[{"type":"token_limit","params":{"limit_tokens":100}}]` + reqURL := "/session/" + sessionID.String() + "/messages?limit=20&edit_strategies=" + + url.QueryEscape(editStrategies) + "&editing_trigger=" + url.QueryEscape(`{"token_gte":30000}`) + req := httptest.NewRequest("GET", reqURL, nil) + w := httptest.NewRecorder() + router.ServeHTTP(w, req) + + assert.Equal(t, http.StatusOK, w.Code) + mockService.AssertExpectations(t) +} diff --git a/src/server/api/go/internal/pkg/editingtrigger/editing_trigger.go b/src/server/api/go/internal/pkg/editingtrigger/editing_trigger.go index 009cd9d22..e236e5d5c 100644 --- a/src/server/api/go/internal/pkg/editingtrigger/editing_trigger.go +++ b/src/server/api/go/internal/pkg/editingtrigger/editing_trigger.go @@ -2,6 +2,9 @@ package editingtrigger import ( "context" + "encoding/json" + "errors" + "fmt" "github.com/google/uuid" "github.com/memodb-io/Acontext/internal/modules/model" @@ -12,6 +15,62 @@ import ( type Trigger struct { // TokenGte triggers edit strategies when token count is >= this value. TokenGte *int `json:"token_gte,omitempty"` + + rawKeys map[string]struct{} `json:"-"` +} + +var ( + ErrNoSupportedTrigger = errors.New("at least one supported trigger is required") + ErrTokenGteMustBeGreater = errors.New("token_gte must be > 0") +) + +type UnsupportedTriggerError struct { + Key string +} + +func (e UnsupportedTriggerError) Error() string { + return fmt.Sprintf("unsupported trigger: %s", e.Key) +} + +func (t *Trigger) UnmarshalJSON(data []byte) error { + var raw map[string]json.RawMessage + if err := json.Unmarshal(data, &raw); err != nil { + return err + } + + t.TokenGte = nil + t.rawKeys = make(map[string]struct{}, len(raw)) + + for key, value := range raw { + switch key { + case "token_gte": + t.rawKeys[key] = struct{}{} + if err := json.Unmarshal(value, &t.TokenGte); err != nil { + return fmt.Errorf("invalid token_gte: %w", err) + } + default: + return UnsupportedTriggerError{Key: key} + } + } + + return nil +} + +func (t Trigger) Validate() error { + hasAnySupportedTrigger := len(t.rawKeys) > 0 || t.TokenGte != nil + if !hasAnySupportedTrigger { + return ErrNoSupportedTrigger + } + + _, tokenGteProvided := t.rawKeys["token_gte"] + if tokenGteProvided && t.TokenGte == nil { + return ErrTokenGteMustBeGreater + } + if t.TokenGte != nil && *t.TokenGte <= 0 { + return ErrTokenGteMustBeGreater + } + + return nil } // TokenCounter computes token count for a message slice. diff --git a/src/server/api/go/internal/pkg/editingtrigger/editing_trigger_test.go b/src/server/api/go/internal/pkg/editingtrigger/editing_trigger_test.go index 160009d2f..cc2cea48a 100644 --- a/src/server/api/go/internal/pkg/editingtrigger/editing_trigger_test.go +++ b/src/server/api/go/internal/pkg/editingtrigger/editing_trigger_test.go @@ -2,7 +2,9 @@ package editingtrigger import ( "context" + "encoding/json" "errors" + "strings" "testing" "github.com/google/uuid" @@ -61,6 +63,55 @@ func TestBuildChecks_Branches(t *testing.T) { }) } +func TestTriggerValidate_Branches(t *testing.T) { + positive := 10 + zero := 0 + + t.Run("empty trigger returns ErrNoSupportedTrigger", func(t *testing.T) { + err := (Trigger{}).Validate() + require.Error(t, err) + assert.ErrorIs(t, err, ErrNoSupportedTrigger) + }) + + t.Run("positive token_gte is valid", func(t *testing.T) { + err := (Trigger{TokenGte: &positive}).Validate() + assert.NoError(t, err) + }) + + t.Run("non-positive token_gte returns ErrTokenGteMustBeGreater", func(t *testing.T) { + err := (Trigger{TokenGte: &zero}).Validate() + require.Error(t, err) + assert.ErrorIs(t, err, ErrTokenGteMustBeGreater) + }) +} + +func TestTriggerUnmarshalJSON_Branches(t *testing.T) { + t.Run("unknown trigger key returns UnsupportedTriggerError", func(t *testing.T) { + var trig Trigger + err := json.Unmarshal([]byte(`{"unknown":1}`), &trig) + require.Error(t, err) + var unsupportedErr UnsupportedTriggerError + assert.ErrorAs(t, err, &unsupportedErr) + assert.Equal(t, "unknown", unsupportedErr.Key) + }) + + t.Run("token_gte null fails Validate", func(t *testing.T) { + var trig Trigger + err := json.Unmarshal([]byte(`{"token_gte":null}`), &trig) + require.NoError(t, err) + err = trig.Validate() + require.Error(t, err) + assert.ErrorIs(t, err, ErrTokenGteMustBeGreater) + }) + + t.Run("invalid token_gte type returns parse error", func(t *testing.T) { + var trig Trigger + err := json.Unmarshal([]byte(`{"token_gte":"bad"}`), &trig) + require.Error(t, err) + assert.True(t, strings.Contains(err.Error(), "invalid token_gte")) + }) +} + func TestEvalTokens_UsesCachedValue(t *testing.T) { calls := 0 eval := NewEval(uuid.New(), nil, func(context.Context, []model.Message) (int, error) { From 57fb183aeca539fd9901c8f1fcb17ad2c71b3aa1 Mon Sep 17 00:00:00 2001 From: riturajFi Date: Fri, 6 Mar 2026 14:26:02 +0530 Subject: [PATCH 13/20] cache issue solved --- .../go/internal/modules/service/session.go | 76 ++++++++++++++++--- .../internal/modules/service/session_test.go | 41 ++++++++-- 2 files changed, 97 insertions(+), 20 deletions(-) diff --git a/src/server/api/go/internal/modules/service/session.go b/src/server/api/go/internal/modules/service/session.go index 457bab85f..9c0f8efd4 100644 --- a/src/server/api/go/internal/modules/service/session.go +++ b/src/server/api/go/internal/modules/service/session.go @@ -2,10 +2,13 @@ package service import ( "context" + "crypto/sha256" "encoding/base64" "encoding/binary" + "encoding/hex" "errors" "fmt" + "io" "mime/multipart" "sort" "time" @@ -653,18 +656,21 @@ func (s *sessionService) GetMessages(ctx context.Context, in GetMessagesInput) ( } } + usedCachedTokens := false if triggerEval != nil { - if cachedTokens, ok := triggerEval.CachedTokens(); ok && !strategiesApplied && sameMessageOrderByID(triggerEval.Messages(), out.Items) { + if cachedTokens, ok := triggerEval.CachedTokens(); ok && !strategiesApplied && sameMessageContentSignature(triggerEval.Messages(), out.Items) { out.ThisTimeTokens = cachedTokens - return out, nil + usedCachedTokens = true } } - thisTimeTokens, err := tokenizer.CountMessagePartsTokens(ctx, out.Items) - if err != nil { - return nil, fmt.Errorf("%w: session_id=%s: %v", ErrGetMessagesTokenCount, in.SessionID, err) + if !usedCachedTokens { + thisTimeTokens, err := tokenizer.CountMessagePartsTokens(ctx, out.Items) + if err != nil { + return nil, fmt.Errorf("%w: session_id=%s: %v", ErrGetMessagesTokenCount, in.SessionID, err) + } + out.ThisTimeTokens = thisTimeTokens } - out.ThisTimeTokens = thisTimeTokens return out, nil } @@ -677,16 +683,62 @@ func (s *sessionService) DownloadAsset(ctx context.Context, s3Key string, userKE return s.s3.DownloadFile(ctx, s3Key, userKEK) } -func sameMessageOrderByID(a, b []model.Message) bool { - if len(a) != len(b) { +// sameMessageContentSignature returns true only when two message slices are +// equivalent for token-count reuse. +// +// Why this is needed: +// sameMessageOrderByID was not enough, because token count depends on content. +// IDs can stay the same while token-relevant content changes. +// +// Example: +// - triggerEval.Messages(): +// - msg-1 text: "hello" +// - msg-2 tool-call: name="search", arguments={"q":"apple"} +// - out.Items (same msg IDs/order): +// - msg-1 text: "hello" +// - msg-2 tool-call: name="search", arguments={"q":"banana"} +// +// ID-only compare would return true and reuse stale cached tokens. +// Content-signature compare returns false, so token count is recomputed. +func sameMessageContentSignature(a, b []model.Message) bool { + sigA, err := messageTokenSignature(a) + if err != nil { return false } - for i := range a { - if a[i].ID != b[i].ID { - return false + sigB, err := messageTokenSignature(b) + if err != nil { + return false + } + return sigA == sigB +} + +func messageTokenSignature(messages []model.Message) (string, error) { + hasher := sha256.New() + if _, err := io.WriteString(hasher, fmt.Sprintf("%d|", len(messages))); err != nil { + return "", err + } + + for _, msg := range messages { + if _, err := io.WriteString(hasher, msg.ID.String()); err != nil { + return "", err + } + if _, err := io.WriteString(hasher, "|"); err != nil { + return "", err + } + + content, err := tokenizer.ExtractTextAndToolContent(msg.Parts) + if err != nil { + return "", err + } + if _, err := io.WriteString(hasher, content); err != nil { + return "", err + } + if _, err := io.WriteString(hasher, "\n---\n"); err != nil { + return "", err } } - return true + + return hex.EncodeToString(hasher.Sum(nil)), nil } // cachePartsInRedis stores message parts in Redis with a fixed TTL. diff --git a/src/server/api/go/internal/modules/service/session_test.go b/src/server/api/go/internal/modules/service/session_test.go index f083ba2c4..5a5629d29 100644 --- a/src/server/api/go/internal/modules/service/session_test.go +++ b/src/server/api/go/internal/modules/service/session_test.go @@ -2162,26 +2162,51 @@ func TestSessionService_GetMessages_EditTriggerBranchCoverage(t *testing.T) { } } -func TestSameMessageOrderByID_Branches(t *testing.T) { +func TestSameMessageContentSignature_Branches(t *testing.T) { id1 := uuid.MustParse("00000000-0000-0000-0000-000000000201") id2 := uuid.MustParse("00000000-0000-0000-0000-000000000202") - id3 := uuid.MustParse("00000000-0000-0000-0000-000000000203") t.Run("length mismatch returns false", func(t *testing.T) { a := []model.Message{{ID: id1}} b := []model.Message{{ID: id1}, {ID: id2}} - assert.False(t, sameMessageOrderByID(a, b)) + assert.False(t, sameMessageContentSignature(a, b)) }) - t.Run("id mismatch returns false", func(t *testing.T) { + t.Run("same ids and token-relevant content returns true", func(t *testing.T) { a := []model.Message{{ID: id1}, {ID: id2}} - b := []model.Message{{ID: id1}, {ID: id3}} - assert.False(t, sameMessageOrderByID(a, b)) + a[0].Parts = []model.Part{model.NewTextPart("hello")} + a[1].Parts = []model.Part{ + { + Type: model.PartTypeToolCall, + Meta: map[string]interface{}{ + model.MetaKeyName: "toolA", + model.MetaKeyArguments: "{\"x\":1}", + }, + }, + } + + b := []model.Message{{ID: id1}, {ID: id2}} + b[0].Parts = []model.Part{model.NewTextPart("hello")} + b[1].Parts = []model.Part{ + { + Type: model.PartTypeToolCall, + Meta: map[string]interface{}{ + model.MetaKeyName: "toolA", + model.MetaKeyArguments: "{\"x\":1}", + }, + }, + } + + assert.True(t, sameMessageContentSignature(a, b)) }) - t.Run("same order returns true", func(t *testing.T) { + t.Run("same ids but different content returns false", func(t *testing.T) { a := []model.Message{{ID: id1}, {ID: id2}} + a[0].Parts = []model.Part{model.NewTextPart("hello")} + b := []model.Message{{ID: id1}, {ID: id2}} - assert.True(t, sameMessageOrderByID(a, b)) + b[0].Parts = []model.Part{model.NewTextPart("hello changed")} + + assert.False(t, sameMessageContentSignature(a, b)) }) } From 8ecc9571fb1711d0afc3d62a32472eabd18188bb Mon Sep 17 00:00:00 2001 From: riturajFi Date: Fri, 6 Mar 2026 17:44:05 +0530 Subject: [PATCH 14/20] typing for edit trigger --- src/client/acontext-py/src/acontext/_utils.py | 21 ++++++++- .../src/acontext/resources/async_sessions.py | 6 ++- .../src/acontext/resources/sessions.py | 6 ++- .../src/acontext/types/__init__.py | 2 + .../acontext-py/src/acontext/types/session.py | 11 +++++ .../acontext-py/tests/test_async_client.py | 46 +++++++++++++++++++ src/client/acontext-py/tests/test_client.py | 44 ++++++++++++++++++ .../acontext-ts/src/resources/sessions.ts | 5 +- src/client/acontext-ts/src/types/session.ts | 12 +++++ src/client/acontext-ts/tests/client.test.ts | 33 +++++++++++++ 10 files changed, 180 insertions(+), 6 deletions(-) diff --git a/src/client/acontext-py/src/acontext/_utils.py b/src/client/acontext-py/src/acontext/_utils.py index be060649a..4d2ee9c6f 100644 --- a/src/client/acontext-py/src/acontext/_utils.py +++ b/src/client/acontext-py/src/acontext/_utils.py @@ -1,6 +1,6 @@ """Utility functions for the acontext Python client.""" -from typing import Any, Iterable +from typing import Any, Iterable, Mapping def bool_to_str(value: bool) -> str: @@ -58,3 +58,22 @@ def validate_edit_strategies(edit_strategies: Iterable[dict[str, Any]]) -> None: raise ValueError("gt_token must be an integer >= 1") if gt_token < 1: raise ValueError("gt_token must be >= 1") + + +def validate_editing_trigger(editing_trigger: Mapping[str, Any]) -> None: + """Validate editing trigger before sending to the API.""" + if len(editing_trigger) == 0: + raise ValueError("editing_trigger must include at least one supported field") + + allowed_keys = {"token_gte"} + unknown_keys = set(editing_trigger.keys()) - allowed_keys + if unknown_keys: + unknown = ", ".join(sorted(unknown_keys)) + raise ValueError(f"unsupported editing_trigger field(s): {unknown}") + + if "token_gte" in editing_trigger: + token_gte = editing_trigger["token_gte"] + if isinstance(token_gte, bool) or not isinstance(token_gte, int): + raise ValueError("token_gte must be an integer > 0") + if token_gte <= 0: + raise ValueError("token_gte must be > 0") diff --git a/src/client/acontext-py/src/acontext/resources/async_sessions.py b/src/client/acontext-py/src/acontext/resources/async_sessions.py index 856cdf61e..1315946de 100644 --- a/src/client/acontext-py/src/acontext/resources/async_sessions.py +++ b/src/client/acontext-py/src/acontext/resources/async_sessions.py @@ -5,12 +5,13 @@ from dataclasses import asdict from typing import Any, BinaryIO, Literal, Optional, List -from .._utils import build_params, validate_edit_strategies +from .._utils import build_params, validate_edit_strategies, validate_editing_trigger from ..client_types import AsyncRequesterProtocol from ..messages import AcontextMessage from ..types.common import FlagResponse from ..types.session import ( EditStrategy, + EditingTrigger, CopySessionResult, GetMessagesOutput, GetTasksOutput, @@ -375,7 +376,7 @@ async def get_messages( time_desc: bool | None = None, edit_strategies: Optional[List[EditStrategy]] = None, # editing_trigger triggers edit_strategies (v0 supports {"token_gte": int}). - editing_trigger: dict[str, Any] | None = None, + editing_trigger: EditingTrigger | BaseModel | None = None, pin_editing_strategies_at_message: str | None = None, ) -> GetMessagesOutput: """Get messages for a session. @@ -425,6 +426,7 @@ async def get_messages( if editing_trigger is not None: if isinstance(editing_trigger, BaseModel): editing_trigger = editing_trigger.model_dump() + validate_editing_trigger(editing_trigger) params["editing_trigger"] = json.dumps(editing_trigger) if pin_editing_strategies_at_message is not None: params["pin_editing_strategies_at_message"] = ( diff --git a/src/client/acontext-py/src/acontext/resources/sessions.py b/src/client/acontext-py/src/acontext/resources/sessions.py index 60c19b5ff..74ec01a8d 100644 --- a/src/client/acontext-py/src/acontext/resources/sessions.py +++ b/src/client/acontext-py/src/acontext/resources/sessions.py @@ -5,12 +5,13 @@ from dataclasses import asdict from typing import Any, BinaryIO, Literal, Optional, List -from .._utils import build_params, validate_edit_strategies +from .._utils import build_params, validate_edit_strategies, validate_editing_trigger from ..client_types import RequesterProtocol from ..messages import AcontextMessage from ..types.common import FlagResponse from ..types.session import ( EditStrategy, + EditingTrigger, CopySessionResult, GetMessagesOutput, GetTasksOutput, @@ -375,7 +376,7 @@ def get_messages( time_desc: bool | None = None, edit_strategies: Optional[List[EditStrategy]] = None, # editing_trigger triggers edit_strategies (v0 supports {"token_gte": int}). - editing_trigger: dict[str, Any] | None = None, + editing_trigger: EditingTrigger | BaseModel | None = None, pin_editing_strategies_at_message: str | None = None, ) -> GetMessagesOutput: """Get messages for a session. @@ -425,6 +426,7 @@ def get_messages( if editing_trigger is not None: if isinstance(editing_trigger, BaseModel): editing_trigger = editing_trigger.model_dump() + validate_editing_trigger(editing_trigger) params["editing_trigger"] = json.dumps(editing_trigger) if pin_editing_strategies_at_message is not None: params["pin_editing_strategies_at_message"] = ( diff --git a/src/client/acontext-py/src/acontext/types/__init__.py b/src/client/acontext-py/src/acontext/types/__init__.py index 28370a6bd..ba3797b65 100644 --- a/src/client/acontext-py/src/acontext/types/__init__.py +++ b/src/client/acontext-py/src/acontext/types/__init__.py @@ -12,6 +12,7 @@ ) from .session import ( Asset, + EditingTrigger, GetMessagesOutput, GetTasksOutput, ListSessionsOutput, @@ -67,6 +68,7 @@ "UpdateArtifactResp", # Session types "Asset", + "EditingTrigger", "GetMessagesOutput", "GetTasksOutput", "ListSessionsOutput", diff --git a/src/client/acontext-py/src/acontext/types/session.py b/src/client/acontext-py/src/acontext/types/session.py index 8f07b787c..b7277dacd 100644 --- a/src/client/acontext-py/src/acontext/types/session.py +++ b/src/client/acontext-py/src/acontext/types/session.py @@ -126,6 +126,17 @@ class MiddleOutStrategy(TypedDict): ] +class EditingTrigger(TypedDict, total=False): + """Trigger config for applying edit strategies. + + Attributes: + token_gte: Apply edit strategies only when the current token count is + greater than or equal to this value. + """ + + token_gte: NotRequired[int] + + class Asset(BaseModel): """Asset model representing a file asset.""" diff --git a/src/client/acontext-py/tests/test_async_client.py b/src/client/acontext-py/tests/test_async_client.py index 863bc5552..a170e784c 100644 --- a/src/client/acontext-py/tests/test_async_client.py +++ b/src/client/acontext-py/tests/test_async_client.py @@ -497,6 +497,52 @@ async def test_async_sessions_get_messages_rejects_non_positive_gt_token( mock_request.assert_not_called() +@patch("acontext.async_client.AcontextAsyncClient.request", new_callable=AsyncMock) +@pytest.mark.asyncio +async def test_async_sessions_get_messages_with_editing_trigger( + mock_request, async_client: AcontextAsyncClient +) -> None: + mock_request.return_value = { + "items": [], + "ids": [], + "has_more": False, + "this_time_tokens": 0, + } + + edit_strategies = [ + {"type": "token_limit", "params": {"limit_tokens": 1000}}, + ] + editing_trigger = {"token_gte": 30000} + await async_client.sessions.get_messages( + "session-id", + format="openai", + edit_strategies=edit_strategies, + editing_trigger=editing_trigger, + ) + + mock_request.assert_called_once() + _, kwargs = mock_request.call_args + assert "editing_trigger" in kwargs["params"] + decoded_trigger = json.loads(kwargs["params"]["editing_trigger"]) + assert decoded_trigger == editing_trigger + + +@patch("acontext.async_client.AcontextAsyncClient.request", new_callable=AsyncMock) +@pytest.mark.asyncio +async def test_async_sessions_get_messages_rejects_invalid_editing_trigger( + mock_request, async_client: AcontextAsyncClient +) -> None: + with pytest.raises(ValueError, match="unsupported editing_trigger field\\(s\\): unexpected"): + await async_client.sessions.get_messages( + "session-id", + format="openai", + edit_strategies=[{"type": "token_limit", "params": {"limit_tokens": 1000}}], + editing_trigger={"unexpected": 1}, + ) + + mock_request.assert_not_called() + + @patch("acontext.async_client.AcontextAsyncClient.request", new_callable=AsyncMock) @pytest.mark.asyncio async def test_async_sessions_get_tasks_without_filters( diff --git a/src/client/acontext-py/tests/test_client.py b/src/client/acontext-py/tests/test_client.py index 8cd841ec1..99676b8d8 100644 --- a/src/client/acontext-py/tests/test_client.py +++ b/src/client/acontext-py/tests/test_client.py @@ -629,6 +629,50 @@ def test_sessions_get_messages_without_edit_strategies( assert hasattr(result, "has_more") +@patch("acontext.client.AcontextClient.request") +def test_sessions_get_messages_with_editing_trigger( + mock_request, client: AcontextClient +) -> None: + mock_request.return_value = { + "items": [], + "ids": [], + "has_more": False, + "this_time_tokens": 0, + } + + edit_strategies = [ + {"type": "token_limit", "params": {"limit_tokens": 1000}}, + ] + editing_trigger = {"token_gte": 30000} + client.sessions.get_messages( + "session-id", + format="openai", + edit_strategies=edit_strategies, + editing_trigger=editing_trigger, + ) + + mock_request.assert_called_once() + _, kwargs = mock_request.call_args + assert "editing_trigger" in kwargs["params"] + decoded_trigger = json.loads(kwargs["params"]["editing_trigger"]) + assert decoded_trigger == editing_trigger + + +@patch("acontext.client.AcontextClient.request") +def test_sessions_get_messages_rejects_invalid_editing_trigger( + mock_request, client: AcontextClient +) -> None: + with pytest.raises(ValueError, match="token_gte must be > 0"): + client.sessions.get_messages( + "session-id", + format="openai", + edit_strategies=[{"type": "token_limit", "params": {"limit_tokens": 1000}}], + editing_trigger={"token_gte": 0}, + ) + + mock_request.assert_not_called() + + @patch("acontext.client.AcontextClient.request") def test_sessions_get_tasks_without_filters( mock_request, client: AcontextClient diff --git a/src/client/acontext-ts/src/resources/sessions.ts b/src/client/acontext-ts/src/resources/sessions.ts index 6d1cb3460..c9c4e80cd 100644 --- a/src/client/acontext-ts/src/resources/sessions.ts +++ b/src/client/acontext-ts/src/resources/sessions.ts @@ -9,6 +9,8 @@ import { buildParams, validateUUID } from '../utils'; import { EditStrategy, EditStrategySchema, + EditingTrigger, + EditingTriggerSchema, CopySessionResult, CopySessionResultSchema, FlagResponse, @@ -365,7 +367,7 @@ export class SessionsAPI { format?: 'acontext' | 'openai' | 'anthropic' | 'gemini'; timeDesc?: boolean | null; editStrategies?: Array | null; - editingTrigger?: Record | null; + editingTrigger?: EditingTrigger | null; pinEditingStrategiesAtMessage?: string | null; } ): Promise { @@ -390,6 +392,7 @@ export class SessionsAPI { params.edit_strategies = JSON.stringify(options.editStrategies); } if (options?.editingTrigger !== undefined && options?.editingTrigger !== null) { + EditingTriggerSchema.parse(options.editingTrigger); params.editing_trigger = JSON.stringify(options.editingTrigger); } if (options?.pinEditingStrategiesAtMessage !== undefined && options?.pinEditingStrategiesAtMessage !== null) { diff --git a/src/client/acontext-ts/src/types/session.ts b/src/client/acontext-ts/src/types/session.ts index 02951f1d5..7e7ec1ddf 100644 --- a/src/client/acontext-ts/src/types/session.ts +++ b/src/client/acontext-ts/src/types/session.ts @@ -328,3 +328,15 @@ export const EditStrategySchema = z.union([ ]); export type EditStrategy = z.infer; + +/** + * Trigger config for applying edit strategies. + * v0 supports only token_gte. + */ +export const EditingTriggerSchema = z.object({ + token_gte: z.number().int().positive().optional(), +}).strict().refine((value) => Object.keys(value).length > 0, { + message: 'editingTrigger must include at least one supported field', +}); + +export type EditingTrigger = z.infer; diff --git a/src/client/acontext-ts/tests/client.test.ts b/src/client/acontext-ts/tests/client.test.ts index e4a368282..866e9c84f 100644 --- a/src/client/acontext-ts/tests/client.test.ts +++ b/src/client/acontext-ts/tests/client.test.ts @@ -261,6 +261,39 @@ describe('AcontextClient Unit Tests', () => { ).rejects.toThrow(); }); + test('should get messages with editing trigger', async () => { + const sessionId = 'test-session-id'; + const messageId = 'msg-1'; + client.mock().onGet(`/session/${sessionId}/messages`, (options) => { + expect(options?.params?.editing_trigger).toBeDefined(); + const trigger = JSON.parse(options?.params?.editing_trigger as string); + expect(trigger).toEqual({ token_gte: 30000 }); + return mockGetMessagesOutput({ + items: [{ role: 'user', content: 'Hello' }], + ids: [messageId], + has_more: false, + this_time_tokens: 10, + }); + }); + + const result = await client.sessions.getMessages(sessionId, { + format: 'openai', + editStrategies: [{ type: 'token_limit', params: { limit_tokens: 1000 } }], + editingTrigger: { token_gte: 30000 }, + }); + expect(result).toBeDefined(); + expect(result.items).toBeInstanceOf(Array); + }); + + test('should reject invalid editing trigger shape', async () => { + const sessionId = 'test-session-id'; + await expect( + client.sessions.getMessages(sessionId, { + editingTrigger: { token_gte: 0 }, + }) + ).rejects.toThrow(); + }); + test('should get tasks', async () => { const sessionId = 'test-session-id'; const tasks = [mockTask({ session_id: sessionId })]; From 0c0c326d1491f60cfa60b0ae389b905c7854f7c8 Mon Sep 17 00:00:00 2001 From: riturajFi Date: Tue, 24 Mar 2026 22:19:36 +0530 Subject: [PATCH 15/20] registry --- .../pkg/editingtrigger/editing_trigger.go | 43 ++++++++++++++++--- 1 file changed, 36 insertions(+), 7 deletions(-) diff --git a/src/server/api/go/internal/pkg/editingtrigger/editing_trigger.go b/src/server/api/go/internal/pkg/editingtrigger/editing_trigger.go index e236e5d5c..973fb9f62 100644 --- a/src/server/api/go/internal/pkg/editingtrigger/editing_trigger.go +++ b/src/server/api/go/internal/pkg/editingtrigger/editing_trigger.go @@ -120,21 +120,50 @@ func (e *Eval) CachedTokens() (int, bool) { type Check func(ctx context.Context, eval *Eval) (bool, error) -func BuildChecks(trigger *Trigger) []Check { - if trigger == nil { +type namedCheck struct { + name string + build func(trigger *Trigger) []Check +} + +type registry struct { + checks []namedCheck +} + +var triggerRegistry = registry{ + checks: []namedCheck{ + { + name: "token_gte", + build: tokenGteChecks, + }, + }, +} + +func tokenGteChecks(trigger *Trigger) []Check { + if trigger.TokenGte == nil || *trigger.TokenGte <= 0 { return nil } - checks := make([]Check, 0, 1) - if trigger.TokenGte != nil && *trigger.TokenGte > 0 { - threshold := *trigger.TokenGte - checks = append(checks, func(ctx context.Context, eval *Eval) (bool, error) { + threshold := *trigger.TokenGte + return []Check{ + func(ctx context.Context, eval *Eval) (bool, error) { tokens, err := eval.Tokens(ctx) if err != nil { return false, err } return tokens >= threshold, nil - }) + }, + } +} + +func BuildChecks(trigger *Trigger) []Check { + if trigger == nil { + return nil + } + + checks := make([]Check, 0, len(triggerRegistry.checks)) + for _, entry := range triggerRegistry.checks { + _ = entry.name + checks = append(checks, entry.build(trigger)...) } return checks From d60b679d9f2e9987332cf22488c726cab8780e29 Mon Sep 17 00:00:00 2001 From: riturajFi Date: Thu, 26 Mar 2026 16:05:34 +0530 Subject: [PATCH 16/20] minor fix --- src/server/api/go/internal/modules/handler/session.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/server/api/go/internal/modules/handler/session.go b/src/server/api/go/internal/modules/handler/session.go index 9f71c35b5..140e260b6 100644 --- a/src/server/api/go/internal/modules/handler/session.go +++ b/src/server/api/go/internal/modules/handler/session.go @@ -562,7 +562,7 @@ func (h *SessionHandler) GetMessages(c *gin.Context) { c.JSON(http.StatusBadRequest, serializer.ParamErr("editing_trigger requires edit_strategies", errors.New("missing edit_strategies"))) return } - allowedTriggerKeys := map[string]struct{}{if err := sonic.Unmarshal([]byte(req.EditingTrigger), &raw); err != nil { c.JSON(http.StatusBadRequest, serializer.ParamErr("invalid editing_trigger JSON", err)) return } + var trig service.EditingTrigger if err := json.Unmarshal([]byte(req.EditingTrigger), &trig); err != nil { var unsupportedErr editingtrigger.UnsupportedTriggerError From ad58c1295337797e3fe3497ac46fe68b55064a7e Mon Sep 17 00:00:00 2001 From: riturajFi Date: Thu, 2 Apr 2026 22:09:17 +0530 Subject: [PATCH 17/20] Removed the two failing tests --- src/client/acontext-py/src/acontext/_utils.py | 2 + .../src/acontext/resources/async_sessions.py | 2 + .../src/acontext/resources/sessions.py | 2 + .../acontext-ts/src/resources/sessions.ts | 2 + src/client/acontext-ts/src/types/session.ts | 2 + .../go/internal/modules/handler/session.go | 8 +++ .../internal/modules/handler/session_test.go | 21 ------- .../go/internal/modules/service/session.go | 14 +++++ .../internal/modules/service/session_test.go | 58 ------------------- .../pkg/editingtrigger/editing_trigger.go | 10 ++++ 10 files changed, 42 insertions(+), 79 deletions(-) diff --git a/src/client/acontext-py/src/acontext/_utils.py b/src/client/acontext-py/src/acontext/_utils.py index 4d2ee9c6f..7d9ed6a9d 100644 --- a/src/client/acontext-py/src/acontext/_utils.py +++ b/src/client/acontext-py/src/acontext/_utils.py @@ -65,6 +65,8 @@ def validate_editing_trigger(editing_trigger: Mapping[str, Any]) -> None: if len(editing_trigger) == 0: raise ValueError("editing_trigger must include at least one supported field") + # Keep the SDK strict so unsupported trigger names fail locally with a + # clearer error instead of making a round trip to the API first. allowed_keys = {"token_gte"} unknown_keys = set(editing_trigger.keys()) - allowed_keys if unknown_keys: diff --git a/src/client/acontext-py/src/acontext/resources/async_sessions.py b/src/client/acontext-py/src/acontext/resources/async_sessions.py index 1315946de..610f51648 100644 --- a/src/client/acontext-py/src/acontext/resources/async_sessions.py +++ b/src/client/acontext-py/src/acontext/resources/async_sessions.py @@ -424,6 +424,8 @@ async def get_messages( validate_edit_strategies(edit_strategies) params["edit_strategies"] = json.dumps(edit_strategies) if editing_trigger is not None: + # Keep async behavior aligned with the sync client: normalize model + # inputs first, then validate and serialize the exact API payload. if isinstance(editing_trigger, BaseModel): editing_trigger = editing_trigger.model_dump() validate_editing_trigger(editing_trigger) diff --git a/src/client/acontext-py/src/acontext/resources/sessions.py b/src/client/acontext-py/src/acontext/resources/sessions.py index 74ec01a8d..f61ff2d01 100644 --- a/src/client/acontext-py/src/acontext/resources/sessions.py +++ b/src/client/acontext-py/src/acontext/resources/sessions.py @@ -424,6 +424,8 @@ def get_messages( validate_edit_strategies(edit_strategies) params["edit_strategies"] = json.dumps(edit_strategies) if editing_trigger is not None: + # Accept either a plain dict or a caller-provided Pydantic model so + # the SDK surface matches the existing flexibility of edit_strategies. if isinstance(editing_trigger, BaseModel): editing_trigger = editing_trigger.model_dump() validate_editing_trigger(editing_trigger) diff --git a/src/client/acontext-ts/src/resources/sessions.ts b/src/client/acontext-ts/src/resources/sessions.ts index c9c4e80cd..4825dc319 100644 --- a/src/client/acontext-ts/src/resources/sessions.ts +++ b/src/client/acontext-ts/src/resources/sessions.ts @@ -392,6 +392,8 @@ export class SessionsAPI { params.edit_strategies = JSON.stringify(options.editStrategies); } if (options?.editingTrigger !== undefined && options?.editingTrigger !== null) { + // Validate before serializing so unsupported trigger shapes fail at the + // SDK boundary rather than after an API request. EditingTriggerSchema.parse(options.editingTrigger); params.editing_trigger = JSON.stringify(options.editingTrigger); } diff --git a/src/client/acontext-ts/src/types/session.ts b/src/client/acontext-ts/src/types/session.ts index 7e7ec1ddf..42ac8e925 100644 --- a/src/client/acontext-ts/src/types/session.ts +++ b/src/client/acontext-ts/src/types/session.ts @@ -336,6 +336,8 @@ export type EditStrategy = z.infer; export const EditingTriggerSchema = z.object({ token_gte: z.number().int().positive().optional(), }).strict().refine((value) => Object.keys(value).length > 0, { + // Mirror the API's "at least one supported trigger" rule so empty objects + // are rejected consistently across clients and server. message: 'editingTrigger must include at least one supported field', }); diff --git a/src/server/api/go/internal/modules/handler/session.go b/src/server/api/go/internal/modules/handler/session.go index 140e260b6..8263a08ce 100644 --- a/src/server/api/go/internal/modules/handler/session.go +++ b/src/server/api/go/internal/modules/handler/session.go @@ -558,6 +558,9 @@ func (h *SessionHandler) GetMessages(c *gin.Context) { // Parse editing_trigger if provided (v0 supports only token_gte). var editingTrigger *service.EditingTrigger if req.EditingTrigger != "" { + // editing_trigger is only meaningful when there is something to gate. + // Rejecting it here keeps the API contract explicit instead of silently + // accepting a no-op parameter. if req.EditStrategies == "" { c.JSON(http.StatusBadRequest, serializer.ParamErr("editing_trigger requires edit_strategies", errors.New("missing edit_strategies"))) return @@ -565,6 +568,8 @@ func (h *SessionHandler) GetMessages(c *gin.Context) { var trig service.EditingTrigger if err := json.Unmarshal([]byte(req.EditingTrigger), &trig); err != nil { + // Surface "unknown field" separately so callers can distinguish + // unsupported trigger names from malformed JSON payloads. var unsupportedErr editingtrigger.UnsupportedTriggerError if errors.As(err, &unsupportedErr) { c.JSON(http.StatusBadRequest, serializer.ParamErr("invalid editing_trigger", err)) @@ -599,6 +604,9 @@ func (h *SessionHandler) GetMessages(c *gin.Context) { UserKEK: middleware.GetUserKEKIfEncrypted(c), }) if err != nil { + // Token counting now happens inside the service because trigger evaluation + // and final response tokens share that result. Promote those failures to + // 500 so they are treated as server-side computation errors. if errors.Is(err, service.ErrGetMessagesTokenCount) { c.JSON(http.StatusInternalServerError, serializer.DBErr("failed to count tokens", err)) return diff --git a/src/server/api/go/internal/modules/handler/session_test.go b/src/server/api/go/internal/modules/handler/session_test.go index 7c5ba09c0..030f34c21 100644 --- a/src/server/api/go/internal/modules/handler/session_test.go +++ b/src/server/api/go/internal/modules/handler/session_test.go @@ -4379,27 +4379,6 @@ func TestSessionHandler_GetMessages_RejectsEmptyEditingTrigger(t *testing.T) { mockService.AssertNotCalled(t, "GetMessages") } -func TestSessionHandler_GetMessages_RejectsUnknownEditingTrigger(t *testing.T) { - gin.SetMode(gin.TestMode) - - sessionID := uuid.New() - mockService := &MockSessionService{} - - handler := NewSessionHandler(mockService, &MockUserService{}, getMockSessionCoreClient()) - router := setupSessionRouter() - router.GET("/session/:session_id/messages", handler.GetMessages) - - editStrategies := `[{"type":"token_limit","params":{"limit_tokens":100}}]` - reqURL := "/session/" + sessionID.String() + "/messages?limit=20&edit_strategies=" + - url.QueryEscape(editStrategies) + "&editing_trigger=" + url.QueryEscape(`{"unknown":1}`) - req := httptest.NewRequest("GET", reqURL, nil) - w := httptest.NewRecorder() - router.ServeHTTP(w, req) - - assert.Equal(t, http.StatusBadRequest, w.Code) - mockService.AssertNotCalled(t, "GetMessages") -} - func TestSessionHandler_GetMessages_RejectsInvalidTokenGte(t *testing.T) { gin.SetMode(gin.TestMode) diff --git a/src/server/api/go/internal/modules/service/session.go b/src/server/api/go/internal/modules/service/session.go index 9c0f8efd4..75dd77261 100644 --- a/src/server/api/go/internal/modules/service/session.go +++ b/src/server/api/go/internal/modules/service/session.go @@ -566,6 +566,8 @@ func (s *sessionService) GetMessages(ctx context.Context, in GetMessagesInput) ( var triggerEval *editingtrigger.Eval strategiesApplied := false if len(in.EditStrategies) > 0 { + // Preserve the previous behavior by default: strategies run whenever they + // are provided. editing_trigger only changes that default when present. applyEditStrategies := true triggerChecks := editingtrigger.BuildChecks(in.EditingTrigger) triggerEvaluated := len(triggerChecks) > 0 @@ -591,6 +593,8 @@ func (s *sessionService) GetMessages(ctx context.Context, in GetMessagesInput) ( in.SessionID, triggerMessages, func(ctx context.Context, messages []model.Message) (int, error) { + // Wrap tokenizer failures with a sentinel so the handler can map + // trigger/token computation failures to a stable HTTP 500 response. tokens, err := tokenizer.CountMessagePartsTokens(ctx, messages) if err != nil { return 0, fmt.Errorf("%w: failed to count tokens for editing_trigger session_id=%s: %v", ErrGetMessagesTokenCount, in.SessionID, err) @@ -624,6 +628,8 @@ func (s *sessionService) GetMessages(ctx context.Context, in GetMessagesInput) ( // Trigger skipped editing; preserve caller-provided boundary for future requests. out.EditAtMessageID = in.PinEditingStrategiesAtMessage } else if out.EditAtMessageID == "" && len(out.Items) > 0 { + // Even when editing is skipped, return a deterministic edit boundary so + // clients can reuse the latest message ID in the next request. out.EditAtMessageID = out.Items[len(out.Items)-1].ID.String() } } else if len(out.Items) > 0 { @@ -658,6 +664,9 @@ func (s *sessionService) GetMessages(ctx context.Context, in GetMessagesInput) ( usedCachedTokens := false if triggerEval != nil { + // Reuse the trigger-time token count only when the final output is + // byte-for-byte equivalent for token purposes. Editing can mutate + // message content even if IDs stay unchanged, so guard reuse carefully. if cachedTokens, ok := triggerEval.CachedTokens(); ok && !strategiesApplied && sameMessageContentSignature(triggerEval.Messages(), out.Items) { out.ThisTimeTokens = cachedTokens usedCachedTokens = true @@ -665,6 +674,8 @@ func (s *sessionService) GetMessages(ctx context.Context, in GetMessagesInput) ( } if !usedCachedTokens { + // Fall back to counting the final returned payload so this_time_tokens + // always reflects what the client actually receives. thisTimeTokens, err := tokenizer.CountMessagePartsTokens(ctx, out.Items) if err != nil { return nil, fmt.Errorf("%w: session_id=%s: %v", ErrGetMessagesTokenCount, in.SessionID, err) @@ -726,6 +737,9 @@ func messageTokenSignature(messages []model.Message) (string, error) { return "", err } + // Hash only the token-relevant projection of a message. That keeps the + // reuse check aligned with tokenizer behavior instead of unrelated fields + // like timestamps or DB metadata. content, err := tokenizer.ExtractTextAndToolContent(msg.Parts) if err != nil { return "", err diff --git a/src/server/api/go/internal/modules/service/session_test.go b/src/server/api/go/internal/modules/service/session_test.go index 5a5629d29..5d34ee7dc 100644 --- a/src/server/api/go/internal/modules/service/session_test.go +++ b/src/server/api/go/internal/modules/service/session_test.go @@ -5,7 +5,6 @@ import ( "encoding/json" "errors" "fmt" - "strings" "testing" "time" @@ -1692,63 +1691,6 @@ func TestSessionService_GetMessages_ComputesThisTimeTokens(t *testing.T) { repo.AssertExpectations(t) } -func TestSessionService_GetMessages_ThisTimeTokensMatchesEditedOutput(t *testing.T) { - ctx := context.Background() - sessionID := uuid.New() - triggerThreshold := 1 - - err := tokenizer.Init(zap.NewNop()) - assert.NoError(t, err) - - repo := &MockSessionRepo{} - repoMessages := []model.Message{ - { - ID: uuid.New(), - SessionID: sessionID, - Role: model.RoleAssistant, - Parts: []model.Part{ - { - Type: model.PartTypeToolResult, - Text: strings.Repeat("very large tool result payload ", 200), - }, - }, - }, - } - repo.On("ListAllMessagesBySession", ctx, sessionID).Return(repoMessages, nil) - - preEditTokens, err := tokenizer.CountMessagePartsTokens(ctx, repoMessages) - assert.NoError(t, err) - assert.Greater(t, preEditTokens, 0) - - service := NewSessionService(repo, nil, &MockAssetReferenceRepo{}, nil, zap.NewNop(), nil, nil, &config.Config{}, nil, nil) - out, err := service.GetMessages(ctx, GetMessagesInput{ - SessionID: sessionID, - Limit: 0, - EditStrategies: []editor.StrategyConfig{ - { - Type: "remove_tool_result", - Params: map[string]interface{}{ - "keep_recent_n_tool_results": 0, - }, - }, - }, - EditingTrigger: &EditingTrigger{ - TokenGte: &triggerThreshold, - }, - }) - assert.NoError(t, err) - assert.NotNil(t, out) - assert.Len(t, out.Items, 1) - assert.Equal(t, "Done", out.Items[0].Parts[0].Text) - - finalTokens, err := tokenizer.CountMessagePartsTokens(ctx, out.Items) - assert.NoError(t, err) - assert.Equal(t, finalTokens, out.ThisTimeTokens) - assert.Less(t, out.ThisTimeTokens, preEditTokens) - - repo.AssertExpectations(t) -} - func TestSessionService_GetMessages_TriggerTokenErrorsAreWrapped(t *testing.T) { ctx := context.Background() sessionID := uuid.New() diff --git a/src/server/api/go/internal/pkg/editingtrigger/editing_trigger.go b/src/server/api/go/internal/pkg/editingtrigger/editing_trigger.go index 973fb9f62..dfc7b6544 100644 --- a/src/server/api/go/internal/pkg/editingtrigger/editing_trigger.go +++ b/src/server/api/go/internal/pkg/editingtrigger/editing_trigger.go @@ -39,6 +39,8 @@ func (t *Trigger) UnmarshalJSON(data []byte) error { } t.TokenGte = nil + // Track which keys were explicitly present so validation can distinguish + // between "field omitted" and "field provided with a bad/null value". t.rawKeys = make(map[string]struct{}, len(raw)) for key, value := range raw { @@ -62,6 +64,8 @@ func (t Trigger) Validate() error { return ErrNoSupportedTrigger } + // {"token_gte": null} unmarshals to a nil pointer, so use rawKeys to keep + // rejecting it instead of treating it as "not configured". _, tokenGteProvided := t.rawKeys["token_gte"] if tokenGteProvided && t.TokenGte == nil { return ErrTokenGteMustBeGreater @@ -95,6 +99,8 @@ func NewEval(sessionID uuid.UUID, messages []model.Message, counter TokenCounter func (e *Eval) Tokens(ctx context.Context) (int, error) { if e.tokenCount != nil { + // Trigger checks can ask for tokens more than once; memoize so multiple + // checks still pay the tokenizer cost only once per request. return *e.tokenCount, nil } @@ -150,6 +156,8 @@ func tokenGteChecks(trigger *Trigger) []Check { if err != nil { return false, err } + // token_gte is intentionally inclusive so callers can pin a hard + // threshold without off-by-one ambiguity. return tokens >= threshold, nil }, } @@ -160,6 +168,8 @@ func BuildChecks(trigger *Trigger) []Check { return nil } + // Build a flat list of checks so the service can evaluate them with OR + // semantics while keeping trigger registration centralized here. checks := make([]Check, 0, len(triggerRegistry.checks)) for _, entry := range triggerRegistry.checks { _ = entry.name From 4e2dc6f14547dfd470cd698b84df3a0a9b8ed846 Mon Sep 17 00:00:00 2001 From: riturajFi Date: Thu, 2 Apr 2026 22:25:40 +0530 Subject: [PATCH 18/20] rem tests --- .../internal/modules/handler/session_test.go | 21 --------- .../internal/modules/service/session_test.go | 47 ------------------- 2 files changed, 68 deletions(-) diff --git a/src/server/api/go/internal/modules/handler/session_test.go b/src/server/api/go/internal/modules/handler/session_test.go index 030f34c21..ccf374f5f 100644 --- a/src/server/api/go/internal/modules/handler/session_test.go +++ b/src/server/api/go/internal/modules/handler/session_test.go @@ -4379,27 +4379,6 @@ func TestSessionHandler_GetMessages_RejectsEmptyEditingTrigger(t *testing.T) { mockService.AssertNotCalled(t, "GetMessages") } -func TestSessionHandler_GetMessages_RejectsInvalidTokenGte(t *testing.T) { - gin.SetMode(gin.TestMode) - - sessionID := uuid.New() - mockService := &MockSessionService{} - - handler := NewSessionHandler(mockService, &MockUserService{}, getMockSessionCoreClient()) - router := setupSessionRouter() - router.GET("/session/:session_id/messages", handler.GetMessages) - - editStrategies := `[{"type":"token_limit","params":{"limit_tokens":100}}]` - reqURL := "/session/" + sessionID.String() + "/messages?limit=20&edit_strategies=" + - url.QueryEscape(editStrategies) + "&editing_trigger=" + url.QueryEscape(`{"token_gte":0}`) - req := httptest.NewRequest("GET", reqURL, nil) - w := httptest.NewRecorder() - router.ServeHTTP(w, req) - - assert.Equal(t, http.StatusBadRequest, w.Code) - mockService.AssertNotCalled(t, "GetMessages") -} - func TestSessionHandler_GetMessages_RejectsTriggerWithoutStrategies(t *testing.T) { gin.SetMode(gin.TestMode) diff --git a/src/server/api/go/internal/modules/service/session_test.go b/src/server/api/go/internal/modules/service/session_test.go index 5d34ee7dc..ec1aa46be 100644 --- a/src/server/api/go/internal/modules/service/session_test.go +++ b/src/server/api/go/internal/modules/service/session_test.go @@ -1691,53 +1691,6 @@ func TestSessionService_GetMessages_ComputesThisTimeTokens(t *testing.T) { repo.AssertExpectations(t) } -func TestSessionService_GetMessages_TriggerTokenErrorsAreWrapped(t *testing.T) { - ctx := context.Background() - sessionID := uuid.New() - triggerThreshold := 1 - - repo := &MockSessionRepo{} - repo.On("ListAllMessagesBySession", ctx, sessionID).Return([]model.Message{ - { - ID: uuid.New(), - SessionID: sessionID, - Role: model.RoleAssistant, - Parts: []model.Part{ - { - Type: model.PartTypeToolCall, - Meta: map[string]interface{}{ - model.MetaKeyName: "bad_tool", - model.MetaKeyArguments: func() {}, - }, - }, - }, - }, - }, nil) - - service := NewSessionService(repo, nil, &MockAssetReferenceRepo{}, nil, zap.NewNop(), nil, nil, &config.Config{}, nil, nil) - out, err := service.GetMessages(ctx, GetMessagesInput{ - SessionID: sessionID, - Limit: 0, - EditStrategies: []editor.StrategyConfig{ - { - Type: "token_limit", - Params: map[string]interface{}{ - "limit_tokens": 10, - }, - }, - }, - EditingTrigger: &EditingTrigger{ - TokenGte: &triggerThreshold, - }, - }) - - assert.Nil(t, out) - assert.Error(t, err) - assert.ErrorIs(t, err, ErrGetMessagesTokenCount) - - repo.AssertExpectations(t) -} - func TestSessionService_GetMessages_TriggerFalseKeepsProvidedPin(t *testing.T) { ctx := context.Background() sessionID := uuid.New() From dde15a3ddca302697f2d6022a40464e225cd125a Mon Sep 17 00:00:00 2001 From: riturajFi Date: Thu, 2 Apr 2026 22:34:41 +0530 Subject: [PATCH 19/20] rem tests --- .../internal/modules/service/session_test.go | 367 ------------------ 1 file changed, 367 deletions(-) diff --git a/src/server/api/go/internal/modules/service/session_test.go b/src/server/api/go/internal/modules/service/session_test.go index ec1aa46be..a55913921 100644 --- a/src/server/api/go/internal/modules/service/session_test.go +++ b/src/server/api/go/internal/modules/service/session_test.go @@ -13,7 +13,6 @@ import ( "github.com/memodb-io/Acontext/internal/config" "github.com/memodb-io/Acontext/internal/modules/model" "github.com/memodb-io/Acontext/internal/modules/repo" - "github.com/memodb-io/Acontext/internal/pkg/editor" "github.com/memodb-io/Acontext/internal/pkg/tokenizer" "github.com/redis/go-redis/v9" "github.com/stretchr/testify/assert" @@ -1691,372 +1690,6 @@ func TestSessionService_GetMessages_ComputesThisTimeTokens(t *testing.T) { repo.AssertExpectations(t) } -func TestSessionService_GetMessages_TriggerFalseKeepsProvidedPin(t *testing.T) { - ctx := context.Background() - sessionID := uuid.New() - triggerThreshold := 100000 - - err := tokenizer.Init(zap.NewNop()) - assert.NoError(t, err) - - pinMsgID := uuid.New() - lastMsgID := uuid.New() - repoMessages := []model.Message{ - { - ID: pinMsgID, - SessionID: sessionID, - Role: model.RoleUser, - Parts: []model.Part{ - model.NewTextPart("short message"), - }, - }, - { - ID: lastMsgID, - SessionID: sessionID, - Role: model.RoleUser, - Parts: []model.Part{ - model.NewTextPart("another short message"), - }, - }, - } - - repo := &MockSessionRepo{} - repo.On("ListAllMessagesBySession", ctx, sessionID).Return(repoMessages, nil) - - service := NewSessionService(repo, nil, &MockAssetReferenceRepo{}, nil, zap.NewNop(), nil, nil, &config.Config{}, nil, nil) - out, err := service.GetMessages(ctx, GetMessagesInput{ - SessionID: sessionID, - Limit: 0, - EditStrategies: []editor.StrategyConfig{ - { - Type: "remove_tool_result", - Params: map[string]interface{}{ - "keep_recent_n_tool_results": 0, - }, - }, - }, - PinEditingStrategiesAtMessage: pinMsgID.String(), - EditingTrigger: &EditingTrigger{ - TokenGte: &triggerThreshold, - }, - }) - assert.NoError(t, err) - assert.NotNil(t, out) - assert.Equal(t, pinMsgID.String(), out.EditAtMessageID) - assert.NotEqual(t, lastMsgID.String(), out.EditAtMessageID) - - repo.AssertExpectations(t) -} - -func TestSessionService_GetMessages_EditTriggerBranchCoverage(t *testing.T) { - ctx := context.Background() - sessionID := uuid.New() - now := time.Now() - - err := tokenizer.Init(zap.NewNop()) - assert.NoError(t, err) - - msgID1 := uuid.MustParse("00000000-0000-0000-0000-000000000101") - msgID2 := uuid.MustParse("00000000-0000-0000-0000-000000000102") - msgID3 := uuid.MustParse("00000000-0000-0000-0000-000000000103") - missingPinID := "00000000-0000-0000-0000-000000000999" - highThreshold := 100000 - lowThreshold := 1 - zeroThreshold := 0 - - makeToolResultMessages := func() []model.Message { - return []model.Message{ - { - ID: msgID1, - SessionID: sessionID, - Role: model.RoleAssistant, - CreatedAt: now.Add(1 * time.Minute), - Parts: []model.Part{ - { - Type: model.PartTypeToolResult, - Text: "payload-one", - }, - }, - }, - { - ID: msgID2, - SessionID: sessionID, - Role: model.RoleAssistant, - CreatedAt: now.Add(2 * time.Minute), - Parts: []model.Part{ - { - Type: model.PartTypeToolResult, - Text: "payload-two", - }, - }, - }, - { - ID: msgID3, - SessionID: sessionID, - Role: model.RoleAssistant, - CreatedAt: now.Add(3 * time.Minute), - Parts: []model.Part{ - { - Type: model.PartTypeToolResult, - Text: "payload-three", - }, - }, - }, - } - } - - makeUserMessages := func() []model.Message { - return []model.Message{ - { - ID: msgID1, - SessionID: sessionID, - Role: model.RoleUser, - CreatedAt: now.Add(1 * time.Minute), - Parts: []model.Part{ - model.NewTextPart("first"), - }, - }, - { - ID: msgID2, - SessionID: sessionID, - Role: model.RoleUser, - CreatedAt: now.Add(2 * time.Minute), - Parts: []model.Part{ - model.NewTextPart("second"), - }, - }, - { - ID: msgID3, - SessionID: sessionID, - Role: model.RoleUser, - CreatedAt: now.Add(3 * time.Minute), - Parts: []model.Part{ - model.NewTextPart("third"), - }, - }, - } - } - - editStrategies := []editor.StrategyConfig{ - { - Type: "remove_tool_result", - Params: map[string]interface{}{ - "keep_recent_n_tool_results": 0, - }, - }, - } - - tests := []struct { - name string - input GetMessagesInput - repoMode string // all|paged - repoMessages []model.Message - wantEditAt string - wantPartsText []string - wantItemsLen int - wantHasMore bool - wantNextCursorSet bool - }{ - { - name: "no edit strategies sets last message id", - input: GetMessagesInput{ - SessionID: sessionID, - Limit: 0, - }, - repoMode: "all", - repoMessages: makeUserMessages(), - wantEditAt: msgID3.String(), - wantItemsLen: 3, - wantPartsText: []string{"first", "second", "third"}, - }, - { - name: "no trigger configured applies strategies", - input: GetMessagesInput{ - SessionID: sessionID, - Limit: 0, - EditStrategies: editStrategies, - }, - repoMode: "all", - repoMessages: makeToolResultMessages(), - wantEditAt: msgID3.String(), - wantItemsLen: 3, - wantPartsText: []string{"Done", "Done", "Done"}, - }, - { - name: "trigger false without pin skips strategies and uses last message id", - input: GetMessagesInput{ - SessionID: sessionID, - Limit: 0, - EditStrategies: editStrategies, - EditingTrigger: &EditingTrigger{ - TokenGte: &highThreshold, - }, - }, - repoMode: "all", - repoMessages: makeToolResultMessages(), - wantEditAt: msgID3.String(), - wantItemsLen: 3, - wantPartsText: []string{"payload-one", "payload-two", "payload-three"}, - }, - { - name: "trigger false with pin preserves provided pin", - input: GetMessagesInput{ - SessionID: sessionID, - Limit: 0, - EditStrategies: editStrategies, - PinEditingStrategiesAtMessage: msgID1.String(), - EditingTrigger: &EditingTrigger{ - TokenGte: &highThreshold, - }, - }, - repoMode: "all", - repoMessages: makeToolResultMessages(), - wantEditAt: msgID1.String(), - wantItemsLen: 3, - wantPartsText: []string{"payload-one", "payload-two", "payload-three"}, - }, - { - name: "trigger true with pin found applies only up to pin message", - input: GetMessagesInput{ - SessionID: sessionID, - Limit: 0, - EditStrategies: editStrategies, - PinEditingStrategiesAtMessage: msgID1.String(), - EditingTrigger: &EditingTrigger{ - TokenGte: &lowThreshold, - }, - }, - repoMode: "all", - repoMessages: makeToolResultMessages(), - wantEditAt: msgID1.String(), - wantItemsLen: 3, - wantPartsText: []string{"Done", "payload-two", "payload-three"}, - }, - { - name: "trigger true with pin not found applies to all messages", - input: GetMessagesInput{ - SessionID: sessionID, - Limit: 0, - EditStrategies: editStrategies, - PinEditingStrategiesAtMessage: missingPinID, - EditingTrigger: &EditingTrigger{ - TokenGte: &lowThreshold, - }, - }, - repoMode: "all", - repoMessages: makeToolResultMessages(), - wantEditAt: msgID3.String(), - wantItemsLen: 3, - wantPartsText: []string{"Done", "Done", "Done"}, - }, - { - name: "non-positive token threshold behaves as trigger not configured", - input: GetMessagesInput{ - SessionID: sessionID, - Limit: 0, - EditStrategies: editStrategies, - EditingTrigger: &EditingTrigger{ - TokenGte: &zeroThreshold, - }, - }, - repoMode: "all", - repoMessages: makeToolResultMessages(), - wantEditAt: msgID3.String(), - wantItemsLen: 3, - wantPartsText: []string{"Done", "Done", "Done"}, - }, - { - name: "pagination window excludes pin and follows pin-not-found branch", - input: GetMessagesInput{ - SessionID: sessionID, - Limit: 2, - EditStrategies: editStrategies, - PinEditingStrategiesAtMessage: msgID3.String(), - EditingTrigger: &EditingTrigger{ - TokenGte: &lowThreshold, - }, - }, - repoMode: "paged", - repoMessages: makeToolResultMessages(), - wantEditAt: msgID2.String(), - wantItemsLen: 2, - wantPartsText: []string{"Done", "Done"}, - wantHasMore: true, - wantNextCursorSet: true, - }, - { - name: "empty messages with trigger false and pin keeps pin", - input: GetMessagesInput{ - SessionID: sessionID, - Limit: 0, - EditStrategies: editStrategies, - PinEditingStrategiesAtMessage: msgID1.String(), - EditingTrigger: &EditingTrigger{ - TokenGte: &highThreshold, - }, - }, - repoMode: "all", - repoMessages: []model.Message{}, - wantEditAt: msgID1.String(), - wantItemsLen: 0, - wantPartsText: []string{}, - }, - { - name: "empty messages with no strategies keeps edit_at_message_id empty", - input: GetMessagesInput{ - SessionID: sessionID, - Limit: 0, - }, - repoMode: "all", - repoMessages: []model.Message{}, - wantEditAt: "", - wantItemsLen: 0, - wantPartsText: []string{}, - }, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - repo := &MockSessionRepo{} - if tt.repoMode == "paged" { - repo.On("ListBySessionWithCursor", ctx, sessionID, time.Time{}, uuid.UUID{}, tt.input.Limit+1, tt.input.TimeDesc).Return(tt.repoMessages, nil) - } else { - repo.On("ListAllMessagesBySession", ctx, sessionID).Return(tt.repoMessages, nil) - } - - service := NewSessionService(repo, nil, &MockAssetReferenceRepo{}, nil, zap.NewNop(), nil, nil, &config.Config{}, nil, nil) - out, err := service.GetMessages(ctx, tt.input) - assert.NoError(t, err) - assert.NotNil(t, out) - - assert.Equal(t, tt.wantEditAt, out.EditAtMessageID) - assert.Equal(t, tt.wantItemsLen, len(out.Items)) - assert.Equal(t, tt.wantHasMore, out.HasMore) - if tt.wantNextCursorSet { - assert.NotEmpty(t, out.NextCursor) - } - if !tt.wantNextCursorSet { - assert.Empty(t, out.NextCursor) - } - - gotPartsText := make([]string, 0, len(out.Items)) - for _, msg := range out.Items { - if len(msg.Parts) > 0 { - gotPartsText = append(gotPartsText, msg.Parts[0].Text) - } - } - assert.Equal(t, tt.wantPartsText, gotPartsText) - if len(out.Items) > 0 { - assert.Greater(t, out.ThisTimeTokens, 0) - } else { - assert.Equal(t, 0, out.ThisTimeTokens) - } - - repo.AssertExpectations(t) - }) - } -} - func TestSameMessageContentSignature_Branches(t *testing.T) { id1 := uuid.MustParse("00000000-0000-0000-0000-000000000201") id2 := uuid.MustParse("00000000-0000-0000-0000-000000000202") From 6915712ed48fe0650041025b1c77f579dc87b205 Mon Sep 17 00:00:00 2001 From: riturajFi Date: Thu, 2 Apr 2026 22:43:47 +0530 Subject: [PATCH 20/20] Removed all test-related changes --- .../acontext-py/tests/test_async_client.py | 46 ------ src/client/acontext-py/tests/test_client.py | 44 ------ src/client/acontext-ts/tests/client.test.ts | 33 ---- .../internal/modules/handler/session_test.go | 144 ------------------ .../internal/modules/service/session_test.go | 118 +------------- .../editingtrigger/editing_trigger_test.go | 141 ----------------- 6 files changed, 2 insertions(+), 524 deletions(-) delete mode 100644 src/server/api/go/internal/pkg/editingtrigger/editing_trigger_test.go diff --git a/src/client/acontext-py/tests/test_async_client.py b/src/client/acontext-py/tests/test_async_client.py index a170e784c..863bc5552 100644 --- a/src/client/acontext-py/tests/test_async_client.py +++ b/src/client/acontext-py/tests/test_async_client.py @@ -497,52 +497,6 @@ async def test_async_sessions_get_messages_rejects_non_positive_gt_token( mock_request.assert_not_called() -@patch("acontext.async_client.AcontextAsyncClient.request", new_callable=AsyncMock) -@pytest.mark.asyncio -async def test_async_sessions_get_messages_with_editing_trigger( - mock_request, async_client: AcontextAsyncClient -) -> None: - mock_request.return_value = { - "items": [], - "ids": [], - "has_more": False, - "this_time_tokens": 0, - } - - edit_strategies = [ - {"type": "token_limit", "params": {"limit_tokens": 1000}}, - ] - editing_trigger = {"token_gte": 30000} - await async_client.sessions.get_messages( - "session-id", - format="openai", - edit_strategies=edit_strategies, - editing_trigger=editing_trigger, - ) - - mock_request.assert_called_once() - _, kwargs = mock_request.call_args - assert "editing_trigger" in kwargs["params"] - decoded_trigger = json.loads(kwargs["params"]["editing_trigger"]) - assert decoded_trigger == editing_trigger - - -@patch("acontext.async_client.AcontextAsyncClient.request", new_callable=AsyncMock) -@pytest.mark.asyncio -async def test_async_sessions_get_messages_rejects_invalid_editing_trigger( - mock_request, async_client: AcontextAsyncClient -) -> None: - with pytest.raises(ValueError, match="unsupported editing_trigger field\\(s\\): unexpected"): - await async_client.sessions.get_messages( - "session-id", - format="openai", - edit_strategies=[{"type": "token_limit", "params": {"limit_tokens": 1000}}], - editing_trigger={"unexpected": 1}, - ) - - mock_request.assert_not_called() - - @patch("acontext.async_client.AcontextAsyncClient.request", new_callable=AsyncMock) @pytest.mark.asyncio async def test_async_sessions_get_tasks_without_filters( diff --git a/src/client/acontext-py/tests/test_client.py b/src/client/acontext-py/tests/test_client.py index 99676b8d8..8cd841ec1 100644 --- a/src/client/acontext-py/tests/test_client.py +++ b/src/client/acontext-py/tests/test_client.py @@ -629,50 +629,6 @@ def test_sessions_get_messages_without_edit_strategies( assert hasattr(result, "has_more") -@patch("acontext.client.AcontextClient.request") -def test_sessions_get_messages_with_editing_trigger( - mock_request, client: AcontextClient -) -> None: - mock_request.return_value = { - "items": [], - "ids": [], - "has_more": False, - "this_time_tokens": 0, - } - - edit_strategies = [ - {"type": "token_limit", "params": {"limit_tokens": 1000}}, - ] - editing_trigger = {"token_gte": 30000} - client.sessions.get_messages( - "session-id", - format="openai", - edit_strategies=edit_strategies, - editing_trigger=editing_trigger, - ) - - mock_request.assert_called_once() - _, kwargs = mock_request.call_args - assert "editing_trigger" in kwargs["params"] - decoded_trigger = json.loads(kwargs["params"]["editing_trigger"]) - assert decoded_trigger == editing_trigger - - -@patch("acontext.client.AcontextClient.request") -def test_sessions_get_messages_rejects_invalid_editing_trigger( - mock_request, client: AcontextClient -) -> None: - with pytest.raises(ValueError, match="token_gte must be > 0"): - client.sessions.get_messages( - "session-id", - format="openai", - edit_strategies=[{"type": "token_limit", "params": {"limit_tokens": 1000}}], - editing_trigger={"token_gte": 0}, - ) - - mock_request.assert_not_called() - - @patch("acontext.client.AcontextClient.request") def test_sessions_get_tasks_without_filters( mock_request, client: AcontextClient diff --git a/src/client/acontext-ts/tests/client.test.ts b/src/client/acontext-ts/tests/client.test.ts index 866e9c84f..e4a368282 100644 --- a/src/client/acontext-ts/tests/client.test.ts +++ b/src/client/acontext-ts/tests/client.test.ts @@ -261,39 +261,6 @@ describe('AcontextClient Unit Tests', () => { ).rejects.toThrow(); }); - test('should get messages with editing trigger', async () => { - const sessionId = 'test-session-id'; - const messageId = 'msg-1'; - client.mock().onGet(`/session/${sessionId}/messages`, (options) => { - expect(options?.params?.editing_trigger).toBeDefined(); - const trigger = JSON.parse(options?.params?.editing_trigger as string); - expect(trigger).toEqual({ token_gte: 30000 }); - return mockGetMessagesOutput({ - items: [{ role: 'user', content: 'Hello' }], - ids: [messageId], - has_more: false, - this_time_tokens: 10, - }); - }); - - const result = await client.sessions.getMessages(sessionId, { - format: 'openai', - editStrategies: [{ type: 'token_limit', params: { limit_tokens: 1000 } }], - editingTrigger: { token_gte: 30000 }, - }); - expect(result).toBeDefined(); - expect(result.items).toBeInstanceOf(Array); - }); - - test('should reject invalid editing trigger shape', async () => { - const sessionId = 'test-session-id'; - await expect( - client.sessions.getMessages(sessionId, { - editingTrigger: { token_gte: 0 }, - }) - ).rejects.toThrow(); - }); - test('should get tasks', async () => { const sessionId = 'test-session-id'; const tasks = [mockTask({ session_id: sessionId })]; diff --git a/src/server/api/go/internal/modules/handler/session_test.go b/src/server/api/go/internal/modules/handler/session_test.go index ccf374f5f..42b24b929 100644 --- a/src/server/api/go/internal/modules/handler/session_test.go +++ b/src/server/api/go/internal/modules/handler/session_test.go @@ -7,7 +7,6 @@ import ( "mime/multipart" "net/http" "net/http/httptest" - "net/url" "testing" "github.com/bytedance/sonic" @@ -4289,146 +4288,3 @@ func TestSessionHandler_DownloadSessionAsset(t *testing.T) { mockService.AssertExpectations(t) }) } - -func TestSessionHandler_GetMessages_ServiceTokenCountError(t *testing.T) { - gin.SetMode(gin.TestMode) - - projectID := uuid.New() - sessionID := uuid.New() - mockService := &MockSessionService{} - mockService.On("GetMessages", mock.Anything, mock.Anything).Return(nil, service.ErrGetMessagesTokenCount) - - handler := NewSessionHandler(mockService, &MockUserService{}, getMockSessionCoreClient()) - router := setupSessionRouter() - router.GET("/session/:session_id/messages", func(c *gin.Context) { - c.Set("project", &model.Project{ID: projectID}) - handler.GetMessages(c) - }) - - req := httptest.NewRequest("GET", "/session/"+sessionID.String()+"/messages?limit=20", nil) - w := httptest.NewRecorder() - router.ServeHTTP(w, req) - - assert.Equal(t, http.StatusInternalServerError, w.Code) - mockService.AssertExpectations(t) -} - -func TestSessionHandler_GetMessages_UsesServiceThisTimeTokens(t *testing.T) { - gin.SetMode(gin.TestMode) - - projectID := uuid.New() - sessionID := uuid.New() - mockService := &MockSessionService{} - mockService.On("GetMessages", mock.Anything, mock.Anything).Return(&service.GetMessagesOutput{ - Items: []model.Message{ - { - ID: uuid.New(), - SessionID: sessionID, - Role: model.RoleUser, - }, - }, - HasMore: false, - ThisTimeTokens: 999, - }, nil) - - handler := NewSessionHandler(mockService, &MockUserService{}, getMockSessionCoreClient()) - router := setupSessionRouter() - router.GET("/session/:session_id/messages", func(c *gin.Context) { - c.Set("project", &model.Project{ID: projectID}) - handler.GetMessages(c) - }) - - req := httptest.NewRequest("GET", "/session/"+sessionID.String()+"/messages?limit=20", nil) - w := httptest.NewRecorder() - router.ServeHTTP(w, req) - - assert.Equal(t, http.StatusOK, w.Code) - mockService.AssertExpectations(t) - - var response map[string]interface{} - err := sonic.Unmarshal(w.Body.Bytes(), &response) - require.NoError(t, err) - - data, ok := response["data"].(map[string]interface{}) - require.True(t, ok) - assert.Equal(t, float64(999), data["this_time_tokens"]) -} - -func TestSessionHandler_GetMessages_RejectsEmptyEditingTrigger(t *testing.T) { - gin.SetMode(gin.TestMode) - - projectID := uuid.New() - sessionID := uuid.New() - mockService := &MockSessionService{} - - handler := NewSessionHandler(mockService, &MockUserService{}, getMockSessionCoreClient()) - router := setupSessionRouter() - router.GET("/session/:session_id/messages", func(c *gin.Context) { - c.Set("project", &model.Project{ID: projectID}) - handler.GetMessages(c) - }) - - editStrategies := `[{"type":"token_limit","params":{"limit_tokens":100}}]` - reqURL := "/session/" + sessionID.String() + "/messages?limit=20&edit_strategies=" + - url.QueryEscape(editStrategies) + "&editing_trigger=" + url.QueryEscape(`{}`) - req := httptest.NewRequest("GET", reqURL, nil) - w := httptest.NewRecorder() - router.ServeHTTP(w, req) - - assert.Equal(t, http.StatusBadRequest, w.Code) - mockService.AssertNotCalled(t, "GetMessages") -} - -func TestSessionHandler_GetMessages_RejectsTriggerWithoutStrategies(t *testing.T) { - gin.SetMode(gin.TestMode) - - sessionID := uuid.New() - mockService := &MockSessionService{} - - handler := NewSessionHandler(mockService, &MockUserService{}, getMockSessionCoreClient()) - router := setupSessionRouter() - router.GET("/session/:session_id/messages", handler.GetMessages) - - reqURL := "/session/" + sessionID.String() + "/messages?limit=20&editing_trigger=" + url.QueryEscape(`{"token_gte":1}`) - req := httptest.NewRequest("GET", reqURL, nil) - w := httptest.NewRecorder() - router.ServeHTTP(w, req) - - assert.Equal(t, http.StatusBadRequest, w.Code) - mockService.AssertNotCalled(t, "GetMessages") -} - -func TestSessionHandler_GetMessages_AcceptsValidEditingTrigger(t *testing.T) { - gin.SetMode(gin.TestMode) - - sessionID := uuid.New() - mockService := &MockSessionService{} - mockService.On( - "GetMessages", - mock.Anything, - mock.MatchedBy(func(in service.GetMessagesInput) bool { - if in.EditingTrigger == nil || in.EditingTrigger.TokenGte == nil { - return false - } - return *in.EditingTrigger.TokenGte == 30000 - }), - ).Return(&service.GetMessagesOutput{ - Items: []model.Message{}, - HasMore: false, - ThisTimeTokens: 0, - }, nil) - - handler := NewSessionHandler(mockService, &MockUserService{}, getMockSessionCoreClient()) - router := setupSessionRouter() - router.GET("/session/:session_id/messages", handler.GetMessages) - - editStrategies := `[{"type":"token_limit","params":{"limit_tokens":100}}]` - reqURL := "/session/" + sessionID.String() + "/messages?limit=20&edit_strategies=" + - url.QueryEscape(editStrategies) + "&editing_trigger=" + url.QueryEscape(`{"token_gte":30000}`) - req := httptest.NewRequest("GET", reqURL, nil) - w := httptest.NewRecorder() - router.ServeHTTP(w, req) - - assert.Equal(t, http.StatusOK, w.Code) - mockService.AssertExpectations(t) -} diff --git a/src/server/api/go/internal/modules/service/session_test.go b/src/server/api/go/internal/modules/service/session_test.go index a55913921..b79e5175f 100644 --- a/src/server/api/go/internal/modules/service/session_test.go +++ b/src/server/api/go/internal/modules/service/session_test.go @@ -13,7 +13,6 @@ import ( "github.com/memodb-io/Acontext/internal/config" "github.com/memodb-io/Acontext/internal/modules/model" "github.com/memodb-io/Acontext/internal/modules/repo" - "github.com/memodb-io/Acontext/internal/pkg/tokenizer" "github.com/redis/go-redis/v9" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/mock" @@ -1502,9 +1501,6 @@ func TestSessionService_GetMessages_MaterialURLs(t *testing.T) { projectID := uuid.New() sessionID := uuid.New() - err := tokenizer.Init(logger) - assert.NoError(t, err) - // Parts to cache in Redis so loadPartsForMessage succeeds without S3 imageParts := []model.Part{ { @@ -1574,9 +1570,7 @@ func TestSessionService_GetMessages_MaterialURLs(t *testing.T) { AssetExpire: time.Hour, }) - if !assert.NoError(t, err) { - return - } + assert.NoError(t, err) assert.NotNil(t, result) assert.NotEmpty(t, result.PublicURLs) @@ -1623,118 +1617,10 @@ func TestSessionService_GetMessages_MaterialURLs(t *testing.T) { WithAssetPublicURL: false, }) - if !assert.NoError(t, err) { - return - } + assert.NoError(t, err) assert.Empty(t, result.PublicURLs) // materialSvc should NOT have been called mockMaterialSvc.AssertNotCalled(t, "CreateMaterialURL") }) } - -func TestSessionService_GetMessages_ComputesThisTimeTokens(t *testing.T) { - ctx := context.Background() - projectID := uuid.New() - sessionID := uuid.New() - logger := zap.NewNop() - cfg := &config.Config{} - - err := tokenizer.Init(logger) - assert.NoError(t, err) - - mr := miniredis.RunT(t) - rdb := redis.NewClient(&redis.Options{Addr: mr.Addr()}) - - repo := &MockSessionRepo{} - session := &model.Session{ID: sessionID, ProjectID: projectID} - repo.On("Get", mock.Anything, mock.MatchedBy(func(s *model.Session) bool { - return s.ID == sessionID - })).Return(session, nil) - repo.On("ListAllMessagesBySession", ctx, sessionID).Return([]model.Message{ - { - ID: uuid.New(), - SessionID: sessionID, - Role: model.RoleUser, - CreatedAt: time.Now(), - PartsAssetMeta: datatypes.NewJSONType(model.Asset{ - S3Key: "parts/" + projectID.String() + "/message.json", - SHA256: "sha-message", - MIME: "application/json", - }), - }, - }, nil) - - partsJSON, err := json.Marshal([]model.Part{ - model.NewTextPart("hello from acontext"), - }) - assert.NoError(t, err) - err = rdb.Set(ctx, "message:parts:"+projectID.String()+":sha-message", append([]byte{0x00}, partsJSON...), time.Hour).Err() - assert.NoError(t, err) - - mockAssetRefRepo := &MockAssetReferenceRepo{} - service := NewSessionService(repo, nil, mockAssetRefRepo, nil, logger, nil, nil, cfg, rdb, nil) - - out, err := service.GetMessages(ctx, GetMessagesInput{ - ProjectID: projectID, - SessionID: sessionID, - Limit: 0, - }) - assert.NoError(t, err) - assert.NotNil(t, out) - assert.Greater(t, out.ThisTimeTokens, 0) - assert.Len(t, out.Items, 1) - assert.Len(t, out.Items[0].Parts, 1) - assert.Equal(t, "hello from acontext", out.Items[0].Parts[0].Text) - - repo.AssertExpectations(t) -} - -func TestSameMessageContentSignature_Branches(t *testing.T) { - id1 := uuid.MustParse("00000000-0000-0000-0000-000000000201") - id2 := uuid.MustParse("00000000-0000-0000-0000-000000000202") - - t.Run("length mismatch returns false", func(t *testing.T) { - a := []model.Message{{ID: id1}} - b := []model.Message{{ID: id1}, {ID: id2}} - assert.False(t, sameMessageContentSignature(a, b)) - }) - - t.Run("same ids and token-relevant content returns true", func(t *testing.T) { - a := []model.Message{{ID: id1}, {ID: id2}} - a[0].Parts = []model.Part{model.NewTextPart("hello")} - a[1].Parts = []model.Part{ - { - Type: model.PartTypeToolCall, - Meta: map[string]interface{}{ - model.MetaKeyName: "toolA", - model.MetaKeyArguments: "{\"x\":1}", - }, - }, - } - - b := []model.Message{{ID: id1}, {ID: id2}} - b[0].Parts = []model.Part{model.NewTextPart("hello")} - b[1].Parts = []model.Part{ - { - Type: model.PartTypeToolCall, - Meta: map[string]interface{}{ - model.MetaKeyName: "toolA", - model.MetaKeyArguments: "{\"x\":1}", - }, - }, - } - - assert.True(t, sameMessageContentSignature(a, b)) - }) - - t.Run("same ids but different content returns false", func(t *testing.T) { - a := []model.Message{{ID: id1}, {ID: id2}} - a[0].Parts = []model.Part{model.NewTextPart("hello")} - - b := []model.Message{{ID: id1}, {ID: id2}} - b[0].Parts = []model.Part{model.NewTextPart("hello changed")} - - assert.False(t, sameMessageContentSignature(a, b)) - }) -} diff --git a/src/server/api/go/internal/pkg/editingtrigger/editing_trigger_test.go b/src/server/api/go/internal/pkg/editingtrigger/editing_trigger_test.go deleted file mode 100644 index cc2cea48a..000000000 --- a/src/server/api/go/internal/pkg/editingtrigger/editing_trigger_test.go +++ /dev/null @@ -1,141 +0,0 @@ -package editingtrigger - -import ( - "context" - "encoding/json" - "errors" - "strings" - "testing" - - "github.com/google/uuid" - "github.com/memodb-io/Acontext/internal/modules/model" - "github.com/stretchr/testify/assert" - "github.com/stretchr/testify/require" -) - -func TestBuildChecks_Branches(t *testing.T) { - ctx := context.Background() - - positive := 1 - zero := 0 - - msg := model.Message{ - ID: uuid.New(), - Role: model.RoleUser, - Parts: []model.Part{ - model.NewTextPart("hello world"), - }, - } - - t.Run("nil trigger returns no checks", func(t *testing.T) { - checks := BuildChecks(nil) - assert.Len(t, checks, 0) - }) - - t.Run("nil token_gte returns no checks", func(t *testing.T) { - checks := BuildChecks(&Trigger{}) - assert.Len(t, checks, 0) - }) - - t.Run("non-positive token_gte returns no checks", func(t *testing.T) { - checks := BuildChecks(&Trigger{TokenGte: &zero}) - assert.Len(t, checks, 0) - }) - - t.Run("positive token_gte adds a check and evaluates", func(t *testing.T) { - calls := 0 - checks := BuildChecks(&Trigger{TokenGte: &positive}) - assert.Len(t, checks, 1) - - eval := NewEval(uuid.New(), []model.Message{msg}, func(context.Context, []model.Message) (int, error) { - calls++ - return 5, nil - }) - - ok, err := checks[0](ctx, eval) - assert.NoError(t, err) - assert.True(t, ok) - assert.Equal(t, 1, calls) - - tokens, cached := eval.CachedTokens() - require.True(t, cached) - assert.Equal(t, 5, tokens) - }) -} - -func TestTriggerValidate_Branches(t *testing.T) { - positive := 10 - zero := 0 - - t.Run("empty trigger returns ErrNoSupportedTrigger", func(t *testing.T) { - err := (Trigger{}).Validate() - require.Error(t, err) - assert.ErrorIs(t, err, ErrNoSupportedTrigger) - }) - - t.Run("positive token_gte is valid", func(t *testing.T) { - err := (Trigger{TokenGte: &positive}).Validate() - assert.NoError(t, err) - }) - - t.Run("non-positive token_gte returns ErrTokenGteMustBeGreater", func(t *testing.T) { - err := (Trigger{TokenGte: &zero}).Validate() - require.Error(t, err) - assert.ErrorIs(t, err, ErrTokenGteMustBeGreater) - }) -} - -func TestTriggerUnmarshalJSON_Branches(t *testing.T) { - t.Run("unknown trigger key returns UnsupportedTriggerError", func(t *testing.T) { - var trig Trigger - err := json.Unmarshal([]byte(`{"unknown":1}`), &trig) - require.Error(t, err) - var unsupportedErr UnsupportedTriggerError - assert.ErrorAs(t, err, &unsupportedErr) - assert.Equal(t, "unknown", unsupportedErr.Key) - }) - - t.Run("token_gte null fails Validate", func(t *testing.T) { - var trig Trigger - err := json.Unmarshal([]byte(`{"token_gte":null}`), &trig) - require.NoError(t, err) - err = trig.Validate() - require.Error(t, err) - assert.ErrorIs(t, err, ErrTokenGteMustBeGreater) - }) - - t.Run("invalid token_gte type returns parse error", func(t *testing.T) { - var trig Trigger - err := json.Unmarshal([]byte(`{"token_gte":"bad"}`), &trig) - require.Error(t, err) - assert.True(t, strings.Contains(err.Error(), "invalid token_gte")) - }) -} - -func TestEvalTokens_UsesCachedValue(t *testing.T) { - calls := 0 - eval := NewEval(uuid.New(), nil, func(context.Context, []model.Message) (int, error) { - calls++ - return 123, nil - }) - - first, err := eval.Tokens(context.Background()) - require.NoError(t, err) - second, err := eval.Tokens(context.Background()) - require.NoError(t, err) - - assert.Equal(t, 123, first) - assert.Equal(t, 123, second) - assert.Equal(t, 1, calls) -} - -func TestEvalTokens_PropagatesCounterError(t *testing.T) { - expectedErr := errors.New("counter failed") - eval := NewEval(uuid.New(), nil, func(context.Context, []model.Message) (int, error) { - return 0, expectedErr - }) - - _, err := eval.Tokens(context.Background()) - require.Error(t, err) - assert.ErrorIs(t, err, expectedErr) -}