Просмотр исходного кода

fix(core): honor Codex input limits

Kit Langton 3 недель назад
Родитель
Сommit
b1a61aaf55

+ 1 - 0
packages/ai/src/schema/options.ts

@@ -123,6 +123,7 @@ export const mergeGenerationOptions = (...items: ReadonlyArray<GenerationOptions
 
 export class ModelLimits extends Schema.Class<ModelLimits>("LLM.ModelLimits")({
   context: Schema.optional(Schema.Number),
+  input: Schema.optional(Schema.Number),
   output: Schema.optional(Schema.Number),
 }) {}
 

+ 1 - 1
packages/core/src/aisdk.ts

@@ -332,7 +332,7 @@ function modelFromLanguage(info: ModelV2.Info, language: LanguageModelV3) {
               body: projected.body === undefined ? undefined : { ...projected.body },
               headers: info.headers,
             },
-      limits: { context: info.limit.context, output: info.limit.output },
+      limits: { context: info.limit.context, input: info.limit.input, output: info.limit.output },
       providerOptions,
     },
     body: {

+ 2 - 2
packages/core/src/model-resolver.ts

@@ -81,7 +81,7 @@ const withDefaults = (model: ModelV2.Info, route: AnyRoute) =>
     headers: providerHeaders(model),
     providerOptions: providerOptions(model),
     http: model.body === undefined ? undefined : { body: model.body },
-    limits: { context: model.limit.context, output: model.limit.output },
+    limits: { context: model.limit.context, input: model.limit.input, output: model.limit.output },
   })
 
 const providerHeaders = (model: ModelV2.Info) => {
@@ -206,7 +206,7 @@ export const fromCatalogModel = (
       ...nativeCredentialSettings(specifier, credential),
       headers: resolved.headers,
       body: resolved.body,
-      limits: { context: resolved.limit.context, output: resolved.limit.output },
+      limits: { context: resolved.limit.context, input: resolved.limit.input, output: resolved.limit.output },
     }
     return yield* Effect.try({
       try: () => {

+ 8 - 0
packages/core/src/plugin/provider/openai.ts

@@ -180,6 +180,12 @@ export const OpenAIPlugin = define({
     })
     yield* load()
     yield* ctx.catalog.transform((evt) => {
+      const codex = evt.provider.get(ProviderV2.ID.make("openai-codex"))
+      if (codex)
+        for (const model of codex.models.values())
+          evt.model.update(codex.provider.id, model.id, (draft) => {
+            draft.enabled = false
+          })
       for (const item of evt.provider.list()) {
         if (!ProviderV2.isAISDK(item.provider.package)) continue
         if (ProviderV2.packageName(item.provider.package) !== "@ai-sdk/openai") continue
@@ -205,6 +211,8 @@ export const OpenAIPlugin = define({
             draft.enabled = false
             return
           }
+          const route = codex?.models.get(draft.id)
+          if (route) draft.limit = { ...route.limit }
           draft.cost = []
         })
       }

+ 3 - 1
packages/core/src/session/compaction.ts

@@ -351,10 +351,12 @@ const make = (dependencies: Dependencies) => {
     )
     if (!last) return false
     const output = Math.min(input.model.route.defaults.limits?.output ?? 0, OUTPUT_TOKEN_MAX)
+    const inputLimit = input.model.route.defaults.limits?.input
+    const limit = inputLimit === undefined ? context - (output || config.buffer) : Math.max(0, inputLimit - config.buffer)
     const used =
       last.tokens.input + last.tokens.output + last.tokens.reasoning + last.tokens.cache.read + last.tokens.cache.write
     if (used <= 0) return false
-    return used >= context - (output || config.buffer)
+    return used >= limit
   }
   const compactManual = Effect.fn("SessionCompaction.compactManual")(function* (input: ManualInput) {
     const content = planContent(input.messages, config.tokens)

+ 4 - 2
packages/core/test/model-resolver.test.ts

@@ -17,6 +17,7 @@ interface ModelOptions {
   readonly headers?: ModelV2.Info["headers"]
   readonly body?: ModelV2.Info["body"]
   readonly variants?: ModelV2.Info["variants"]
+  readonly limit?: ModelV2.Info["limit"]
 }
 
 const model = (packageName: string | undefined, options: ModelOptions = {}) =>
@@ -36,7 +37,7 @@ const model = (packageName: string | undefined, options: ModelOptions = {}) =>
     cost: [],
     status: "active",
     enabled: true,
-    limit: { context: 100, output: 20 },
+    limit: options.limit ?? { context: 100, output: 20 },
   })
 
 describe("ModelResolver", () => {
@@ -44,6 +45,7 @@ describe("ModelResolver", () => {
     Effect.gen(function* () {
       const catalog = model(ProviderV2.aisdk("@ai-sdk/openai"), {
         settings: { baseURL: "https://openai.example/v1" },
+        limit: { context: 100, input: 70, output: 20 },
       })
       const resolved = yield* ModelResolver.fromCatalogModel(catalog)
 
@@ -55,7 +57,7 @@ describe("ModelResolver", () => {
         endpoint: { baseURL: "https://openai.example/v1" },
         defaults: {
           headers: { "x-test": "header" },
-          limits: { context: 100, output: 20 },
+          limits: { context: 100, input: 70, output: 20 },
           http: { body: { custom_extension: { enabled: true } } },
         },
       })

+ 36 - 6
packages/core/test/plugin/provider-openai.test.ts

@@ -164,6 +164,16 @@ describe("OpenAIPlugin", () => {
       const catalog = yield* Catalog.Service
       const credentials = yield* Credential.Service
       yield* catalog.transform((catalog) => {
+        const codex = ProviderV2.Info.make({
+          ...ProviderV2.Info.empty(ProviderV2.ID.make("openai-codex")),
+          package: ProviderV2.aisdk("@ai-sdk/openai"),
+        })
+        catalog.provider.update(codex.id, (draft) => {
+          draft.package = codex.package
+        })
+        catalog.model.update(codex.id, ModelV2.ID.make("gpt-5.6-sol"), (model) => {
+          model.limit = { context: 400_000, input: 272_000, output: 128_000 }
+        })
         const item = ProviderV2.Info.make({
           ...ProviderV2.Info.empty(ProviderV2.ID.openai),
           package: ProviderV2.aisdk("@ai-sdk/openai"),
@@ -189,7 +199,9 @@ describe("OpenAIPlugin", () => {
           model.body = { reasoning: { mode: "pro" } }
         })
         catalog.model.update(item.id, ModelV2.ID.make("gpt-5.6"), () => {})
-        catalog.model.update(item.id, ModelV2.ID.make("gpt-5.6-sol"), () => {})
+        catalog.model.update(item.id, ModelV2.ID.make("gpt-5.6-sol"), (model) => {
+          model.limit = { context: 1_050_000, input: 922_000, output: 128_000 }
+        })
         catalog.model.update(item.id, ModelV2.ID.make("gpt-4.1"), () => {})
       })
       yield* credentials.create({
@@ -215,9 +227,13 @@ describe("OpenAIPlugin", () => {
         false,
       )
       expect(required(yield* catalog.model.get(ProviderV2.ID.openai, ModelV2.ID.make("gpt-5.6"))).enabled).toBe(false)
-      expect(required(yield* catalog.model.get(ProviderV2.ID.openai, ModelV2.ID.make("gpt-5.6-sol"))).enabled).toBe(
-        true,
-      )
+      const sol = required(yield* catalog.model.get(ProviderV2.ID.openai, ModelV2.ID.make("gpt-5.6-sol")))
+      expect(sol.enabled).toBe(true)
+      expect(sol.limit).toEqual({ context: 400_000, input: 272_000, output: 128_000 })
+      expect(
+        required(yield* catalog.model.get(ProviderV2.ID.make("openai-codex"), ModelV2.ID.make("gpt-5.6-sol")))
+          .enabled,
+      ).toBe(false)
       expect(required(yield* catalog.model.get(ProviderV2.ID.openai, ModelV2.ID.make("gpt-4.1"))).enabled).toBe(false)
     }),
   )
@@ -227,6 +243,16 @@ describe("OpenAIPlugin", () => {
       const catalog = yield* Catalog.Service
       const credentials = yield* Credential.Service
       yield* catalog.transform((catalog) => {
+        const codex = ProviderV2.Info.make({
+          ...ProviderV2.Info.empty(ProviderV2.ID.make("openai-codex")),
+          package: ProviderV2.aisdk("@ai-sdk/openai"),
+        })
+        catalog.provider.update(codex.id, (draft) => {
+          draft.package = codex.package
+        })
+        catalog.model.update(codex.id, ModelV2.ID.make("gpt-5.5"), (model) => {
+          model.limit = { context: 400_000, input: 272_000, output: 128_000 }
+        })
         const item = ProviderV2.Info.make({
           ...ProviderV2.Info.empty(ProviderV2.ID.openai),
           package: ProviderV2.aisdk("@ai-sdk/openai"),
@@ -234,7 +260,9 @@ describe("OpenAIPlugin", () => {
         catalog.provider.update(item.id, (draft) => {
           draft.package = item.package
         })
-        catalog.model.update(item.id, ModelV2.ID.make("gpt-5.5"), () => {})
+        catalog.model.update(item.id, ModelV2.ID.make("gpt-5.5"), (model) => {
+          model.limit = { context: 1_050_000, input: 922_000, output: 128_000 }
+        })
         catalog.model.update(item.id, ModelV2.ID.make("gpt-4.1"), () => {})
       })
       yield* credentials.create({
@@ -243,7 +271,9 @@ describe("OpenAIPlugin", () => {
       })
       yield* addPlugin()
 
-      expect(required(yield* catalog.model.get(ProviderV2.ID.openai, ModelV2.ID.make("gpt-5.5"))).enabled).toBe(true)
+      const model = required(yield* catalog.model.get(ProviderV2.ID.openai, ModelV2.ID.make("gpt-5.5")))
+      expect(model.enabled).toBe(true)
+      expect(model.limit).toEqual({ context: 1_050_000, input: 922_000, output: 128_000 })
       expect(required(yield* catalog.model.get(ProviderV2.ID.openai, ModelV2.ID.make("gpt-4.1"))).enabled).toBe(true)
     }),
   )

+ 38 - 0
packages/core/test/session-compaction.test.ts

@@ -16,11 +16,15 @@ import { SessionRunnerModel } from "@opencode-ai/core/session/runner/model"
 import { SessionTable } from "@opencode-ai/core/session/sql"
 import { SessionStore } from "@opencode-ai/core/session/store"
 import { SessionV2 } from "@opencode-ai/core/session"
+import { AgentV2 } from "@opencode-ai/core/agent"
+import { Location } from "@opencode-ai/core/location"
 import { Project } from "@opencode-ai/core/project"
 import { ProjectTable } from "@opencode-ai/core/project/sql"
 import { App } from "@opencode-ai/core/app"
 import { AbsolutePath } from "@opencode-ai/core/schema"
 import { Money } from "@opencode-ai/schema/money"
+import { ModelV2 } from "@opencode-ai/core/model"
+import { ProviderV2 } from "@opencode-ai/core/provider"
 import { DateTime, Effect, Fiber, Layer, Stream } from "effect"
 import { asc, eq } from "drizzle-orm"
 import { testEffect } from "./lib/effect"
@@ -131,6 +135,40 @@ test("compaction prompt requires the checkpoint headings in order", () => {
   expect(prompt).toContain("Keep every section, even when empty.")
 })
 
+it.effect("auto compaction uses the model input limit", () =>
+  Effect.gen(function* () {
+    const compaction = yield* SessionCompaction.Service
+    const session = SessionV2.Info.make({
+      id: SessionV2.ID.make("ses_input_limit"),
+      projectID: Project.ID.global,
+      title: "Input limit",
+      cost: Money.USD.zero,
+      tokens: { input: 0, output: 0, reasoning: 0, cache: { read: 0, write: 0 } },
+      time: { created: DateTime.makeUnsafe(0), updated: DateTime.makeUnsafe(0) },
+      location: Location.Ref.make({ directory: AbsolutePath.make("/project") }),
+    })
+    const route = Model.make({
+      id: "codex-model",
+      provider: "openai",
+      route: OpenAIChat.route.with({ limits: { context: 400_000, input: 272_000, output: 128_000 } }),
+    })
+    const message = (input: number) =>
+      SessionMessage.Assistant.make({
+        id: SessionMessage.ID.create(),
+        type: "assistant",
+        agent: AgentV2.defaultID,
+        model: { id: ModelV2.ID.make("codex-model"), providerID: ProviderV2.ID.openai },
+        content: [],
+        cost: Money.USD.zero,
+        tokens: { input, output: 0, reasoning: 0, cache: { read: 0, write: 0 } },
+        time: { created: DateTime.makeUnsafe(0) },
+      })
+
+    expect(compaction.required({ session, model: route, cost: [], messages: [message(251_999)] })).toBe(false)
+    expect(compaction.required({ session, model: route, cost: [], messages: [message(252_000)] })).toBe(true)
+  }),
+)
+
 it.effect("manual compaction summarizes short context instead of no-op", () =>
   Effect.gen(function* () {
     requests = []

+ 2 - 1
packages/tui/src/util/session.ts

@@ -29,9 +29,10 @@ export function contextUsage(
     last.tokens.input + last.tokens.output + last.tokens.reasoning + last.tokens.cache.read + last.tokens.cache.write
   if (tokens <= 0) return
   const model = models?.find((model) => model.providerID === last.model.providerID && model.id === last.model.id)
+  const limit = model?.limit.input ?? model?.limit.context
   return {
     tokens,
-    percent: model?.limit.context ? Math.round((tokens / model.limit.context) * 100) : undefined,
+    percent: limit ? Math.round((tokens / limit) * 100) : undefined,
   }
 }
 

+ 17 - 2
packages/tui/test/util/session.test.ts

@@ -1,6 +1,6 @@
 import { describe, expect, test } from "bun:test"
-import type { SessionMessageInfo } from "@opencode-ai/client"
-import { isDefaultTitle, lastAssistantWithUsage } from "../../src/util/session"
+import type { ModelInfo, SessionMessageInfo } from "@opencode-ai/client"
+import { contextUsage, isDefaultTitle, lastAssistantWithUsage } from "../../src/util/session"
 
 const assistant = (id: string, input: number): SessionMessageInfo => ({
   id,
@@ -45,4 +45,19 @@ describe("util.session", () => {
     messages.push(assistant("msg_after", 5))
     expect(lastAssistantWithUsage(messages)?.tokens.input).toBe(5)
   })
+
+  test("reports usage against the effective input window", () => {
+    const models = [
+      {
+        id: "model",
+        providerID: "provider",
+        limit: { context: 400_000, input: 272_000, output: 128_000 },
+      } as unknown as ModelInfo,
+    ]
+
+    expect(contextUsage([assistant("msg_input_limit", 252_000)], models)).toEqual({
+      tokens: 252_000,
+      percent: 93,
+    })
+  })
 })