| 1 | package provider |
| 2 | |
| 3 | import ( |
| 4 | "context" |
| 5 | "testing" |
| 6 | ) |
| 7 | |
| 8 | type legacySharedProvider struct { |
| 9 | budget int |
| 10 | shared bool |
| 11 | } |
| 12 | |
| 13 | func (legacySharedProvider) Name() string { return "legacy" } |
| 14 | func (legacySharedProvider) Stream(context.Context, Request) (<-chan Chunk, error) { |
| 15 | return nil, nil |
| 16 | } |
| 17 | func (p legacySharedProvider) OutputBudget() int { return p.budget } |
| 18 | func (p legacySharedProvider) SharesContextWindow() bool { return p.shared } |
| 19 | |
| 20 | type policyProvider struct { |
| 21 | policy ContextBudgetPolicy |
| 22 | } |
| 23 | |
| 24 | func (policyProvider) Name() string { return "policy" } |
| 25 | func (policyProvider) Stream(context.Context, Request) (<-chan Chunk, error) { |
| 26 | return nil, nil |
| 27 | } |
| 28 | func (p policyProvider) ContextBudgetPolicy() ContextBudgetPolicy { return p.policy } |
| 29 | |
| 30 | func TestResolveContextBudgetPolicyFallsBackToLegacyInterfaces(t *testing.T) { |
| 31 | got := ResolveContextBudgetPolicy(legacySharedProvider{budget: 128_000, shared: true}) |
| 32 | if got.WindowMode != ContextWindowShared || got.AutoOutputTokens != 128_000 || got.LimitMode != OutputLimitAlways { |
| 33 | t.Fatalf("legacy shared policy = %+v", got) |
| 34 | } |
| 35 | unknown := ResolveContextBudgetPolicy(legacySharedProvider{budget: 0, shared: false}) |
| 36 | if unknown.WindowMode != ContextWindowUnknown || unknown.AutoOutputTokens != 0 { |
| 37 | t.Fatalf("legacy unknown policy = %+v", unknown) |
| 38 | } |
| 39 | } |
| 40 | |
| 41 | func TestResolveContextBudgetPolicyPrefersNewInterface(t *testing.T) { |
| 42 | want := ContextBudgetPolicy{ |
| 43 | WindowMode: ContextWindowShared, |
| 44 | AutoOutputTokens: DeepSeekMaxOutputTokens, |
| 45 | MaxOutputTokens: DeepSeekMaxOutputTokens, |
| 46 | LimitMode: OutputLimitOmitWhenSafe, |
| 47 | } |
| 48 | got := ResolveContextBudgetPolicy(policyProvider{policy: want}) |
| 49 | if got != want { |
| 50 | t.Fatalf("policy = %+v, want %+v", got, want) |
| 51 | } |
| 52 | } |
| 53 | |
| 54 | func TestResolveContextBudgetPolicyNilProvider(t *testing.T) { |
| 55 | if got := ResolveContextBudgetPolicy(nil); got.WindowMode != ContextWindowUnknown { |
| 56 | t.Fatalf("nil policy = %+v", got) |
| 57 | } |
| 58 | } |
| 59 |