feat: implement context usage anchor for improved token estimation

- Introduced `ContextUsageAnchor` struct to track context input tokens and message count for conversations. - Updated token estimation functions to utilize the context usage anchor, enhancing accuracy in estimating tokens for projected messages. - Refactored compaction logic to incorporate context usage anchor, allowing for more efficient management of token budgets during model runs. - Added tests to validate the behavior of the context usage anchor across different scenarios, including model switching and message additions.
2026-10-06 13:44:21 +08:00 · 2026-09-01 16:07:51 +08:00
parent 8c6c415a84
commit d004139526
6 changed files with 394 additions and 21 deletions
@@ -297,6 +297,108 @@ async fn automatic_compaction_preflights_provider_input_and_records_rebuilt_toke
        }));
 }

+#[tokio::test]
+async fn incremental_preflight_uses_conversation_anchor_across_model_switch() {
+    let (_directory, store) = fixtures::temp_store().await;
+    let model_a = store
+        .create_model(&ModelConfigInput {
+            sort_order: 0,
+            display_name: "Anchor Model A".into(),
+            group_name: None,
+            model_type: ModelType::OpenAi,
+            base_url: "https://example.com/v1/chat/completions".into(),
+            use_full_url: true,
+            api_key: "test-key".into(),
+            tooltip_data: "Anchor Model A".into(),
+            model_id: "anchor-model-a".into(),
+            reasoning_effort: None,
+            openai_endpoint: OPENAI_CHAT_ENDPOINT.into(),
+            openai_extra_params_enabled: false,
+            openai_extra_params: serde_json::json!({}),
+            custom_headers_enabled: false,
+            custom_headers: serde_json::json!({}),
+            anthropic_extra_params_enabled: false,
+            anthropic_extra_params: serde_json::json!({}),
+            context_window_tokens: None,
+            max_completion_tokens: None,
+            anthropic_max_tokens: None,
+            anthropic_thinking_effort: None,
+            thinking_budget_tokens: None,
+        })
+        .await
+        .unwrap();
+    let model_b = store
+        .create_model(&ModelConfigInput {
+            sort_order: 1,
+            display_name: "Anchor Model B".into(),
+            group_name: None,
+            model_type: ModelType::OpenAi,
+            base_url: "https://example.com/v1/chat/completions".into(),
+            use_full_url: true,
+            api_key: "test-key".into(),
+            tooltip_data: "Anchor Model B".into(),
+            model_id: "anchor-model-b".into(),
+            reasoning_effort: None,
+            openai_endpoint: OPENAI_CHAT_ENDPOINT.into(),
+            openai_extra_params_enabled: false,
+            openai_extra_params: serde_json::json!({}),
+            custom_headers_enabled: false,
+            custom_headers: serde_json::json!({}),
+            anthropic_extra_params_enabled: false,
+            anthropic_extra_params: serde_json::json!({}),
+            context_window_tokens: Some(200_000),
+            max_completion_tokens: None,
+            anthropic_max_tokens: None,
+            anthropic_thinking_effort: None,
+            thinking_budget_tokens: None,
+        })
+        .await
+        .unwrap();
+    let provider = fake_provider::FakeProvider::default();
+    provider.push(text_response("old answer", 103_904, 12));
+    provider.push(text_response("new answer", 104_000, 12));
+    let assets = PromptAssets::load(
+        std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
+            .join("prompt/cursor")
+            .as_path(),
+    )
+    .unwrap();
+    let registry = TransportRegistry::new(
+        store,
+        Arc::new(provider.clone()),
+        PromptCompiler::new(assets),
+    );
+
+    let first = run(
+        &registry,
+        "anchor-first",
+        user_request(
+            "anchor-conversation",
+            "anchor-user-1",
+            &"x".repeat(400_000),
+            &model_a.model_hash,
+            None,
+        ),
+    )
+    .await;
+    let second = run(
+        &registry,
+        "anchor-second",
+        user_request(
+            "anchor-conversation",
+            "anchor-user-2",
+            "short follow-up",
+            &model_b.model_hash,
+            first.checkpoints.last().cloned(),
+        ),
+    )
+    .await;
+
+    assert_eq!(second.summary_started, 0);
+    assert_eq!(second.summary_completed, 0);
+    assert_eq!(provider.requests().len(), 2);
+}
+
 #[tokio::test]
 async fn irreducibly_oversized_current_input_fails_before_provider_dispatch() {
    let (_directory, store) = fixtures::temp_store().await;