smashimo/oh-my-pi
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(coding-agent): added Ollama context discovery and attribution control

Browse Source 
- Added automatic Ollama model context window discovery from metadata for accurate token limits.
- Added `attribution` option to `PromptOptions` for explicit billing and initiator attribution control.
- Added automatic clearing of completed and abandoned todo tasks after ~1 minute.
- Changed session directory migration to use `-tmp-` prefix instead of double-dash format.
- Updated Ollama model registration to use discovered context window instead of hardcoded 128000 token default.

Fixes #440
This commit is contained in:
can1357
2026-03-15 23:05:20 +01:00
parent bdcc08a50c
commit 82b70ead61
5 changed files with 160 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/coding-agent/CHANGELOG.md
+4
View File
@@ -1,13 +1,17 @@
# Changelog
## [Unreleased]
### Added
- Automatic discovery of Ollama model context window from model metadata, enabling accurate token limit configuration
- Added `attribution` option to `PromptOptions` to explicitly control billing/initiator attribution for prompts
- Added automatic clearing of completed and abandoned todo tasks after ~1 minute
### Changed
- Ollama model registration now uses discovered context window instead of hardcoded 128000 token default
- Ollama model maxTokens now respects discovered context window constraints
- Improved session directory migration to handle legacy absolute paths with double-dash format, automatically relocating them to new canonical locations
- Enhanced session directory encoding to use `-tmp-` prefix for temporary directories instead of legacy double-dash format for better clarity
- Updated `SessionManager.create()` to require both `cwd` and `sessionDir` parameters for explicit session directory control
packages/coding-agent/src/commit/agentic/agent.ts
+8 -2
View File
@@ -161,11 +161,17 @@ export async function runCommitAgentSession(input: CommitAgentInput): Promise<Co
let retryCount = 0;
const needsChangelog = input.requireChangelog && input.changelogTargets.length > 0;
await session.prompt(prompt, { expandPromptTemplates: false });
await session.prompt(prompt, {
attribution: "agent",
expandPromptTemplates: false,
});
while (retryCount < MAX_RETRIES && !isProposalComplete(state, needsChangelog)) {
retryCount += 1;
const reminder = buildReminderMessage(state, needsChangelog, retryCount, MAX_RETRIES);
await session.prompt(reminder, { expandPromptTemplates: false });
await session.prompt(reminder, {
attribution: "agent",
expandPromptTemplates: false,
});
}
return state;
packages/coding-agent/src/config/model-registry.ts
+51 -4
View File
@@ -366,6 +366,12 @@ interface CustomModelsResult {
found: boolean;
}
type OllamaDiscoveredModelMetadata = {
reasoning: boolean;
input: ("text" | "image")[];
contextWindow?: number;
};
/**
* Resolve an API key config value to an actual key.
* Checks environment variable first, then treats as literal.
@@ -376,6 +382,40 @@ function resolveApiKeyConfig(keyConfig: string): string | undefined {
return keyConfig;
}
function toPositiveNumberOrUndefined(value: unknown): number | undefined {
if (typeof value === "number" && Number.isFinite(value) && value > 0) {
return value;
}
if (typeof value === "string" && value.trim()) {
const parsed = Number(value);
if (Number.isFinite(parsed) && parsed > 0) {
return parsed;
}
}
return undefined;
}
function extractOllamaContextWindow(payload: Record<string, unknown>): number | undefined {
const modelInfo = payload.model_info;
if (isRecord(modelInfo)) {
for (const [key, value] of Object.entries(modelInfo)) {
if (key === "context_length" || key.endsWith(".context_length")) {
const contextWindow = toPositiveNumberOrUndefined(value);
if (contextWindow !== undefined) {
return contextWindow;
}
}
}
}
const parameters = payload.parameters;
if (typeof parameters !== "string") {
return undefined;
}
const match = parameters.match(/(?:^|\n)\s*num_ctx\s+(\d+)\s*(?:$|\n)/m);
return match ? toPositiveNumberOrUndefined(match[1]) : undefined;
}
function extractGoogleOAuthToken(value: string | undefined): string | undefined {
if (!isAuthenticated(value)) return undefined;
try {
@@ -1096,7 +1136,7 @@ export class ModelRegistry {
endpoint: string,
modelId: string,
headers: Record<string, string> | undefined,
): Promise<{ reasoning: boolean; input: ("text" | "image")[] } | null> {
): Promise<OllamaDiscoveredModelMetadata | null> {
const showUrl = `${endpoint}/api/show`;
try {
const response = await fetch(showUrl, {
@@ -1112,6 +1152,7 @@ export class ModelRegistry {
if (!isRecord(payload)) {
return null;
}
const contextWindow = extractOllamaContextWindow(payload);
const capabilities = payload.capabilities;
if (Array.isArray(capabilities)) {
const normalized = new Set(
@@ -1121,15 +1162,21 @@ export class ModelRegistry {
return {
reasoning: normalized.has("thinking"),
input: supportsVision ? ["text", "image"] : ["text"],
contextWindow,
};
}
if (!isRecord(capabilities)) {
return null;
return {
reasoning: false,
input: ["text"],
contextWindow,
};
}
const supportsVision = capabilities.vision === true || capabilities.image === true;
return {
reasoning: capabilities.thinking === true,
input: supportsVision ? ["text", "image"] : ["text"],
contextWindow,
};
} catch {
return null;
@@ -1170,8 +1217,8 @@ export class ModelRegistry {
reasoning: metadata?.reasoning ?? false,
input: metadata?.input ?? ["text"],
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
contextWindow: 128000,
maxTokens: 8192,
contextWindow: metadata?.contextWindow ?? 128000,
maxTokens: Math.min(metadata?.contextWindow ?? Number.POSITIVE_INFINITY, 8192),
headers: providerConfig.headers,
});
});
packages/coding-agent/test/commit-agentic-attribution.test.ts
+59
View File
@@ -0,0 +1,59 @@
import { afterEach, describe, expect, it, vi } from "bun:test";
import { getBundledModel } from "@oh-my-pi/pi-ai";
import { runCommitAgentSession } from "../src/commit/agentic/agent";
import * as toolsModule from "../src/commit/agentic/tools";
import { Settings } from "../src/config/settings";
import * as sdkModule from "../src/sdk";
import type { PromptOptions } from "../src/session/agent-session";
vi.mock("../src/sdk", () => ({
createAgentSession: vi.fn(),
}));
vi.mock("../src/commit/agentic/tools", () => ({
createCommitTools: vi.fn(() => []),
}));
describe("commit agent prompt attribution", () => {
afterEach(() => {
vi.restoreAllMocks();
});
it("marks generated commit prompts and reminders as agent-attributed", async () => {
const prompts: Array<{ text: string; options?: PromptOptions }> = [];
const session = {
prompt: async (text: string, options?: PromptOptions) => {
prompts.push({ text, options });
},
subscribe: () => () => {},
dispose: async () => {},
};
(sdkModule.createAgentSession as unknown as { mockResolvedValue: (value: unknown) => void }).mockResolvedValue({
session,
});
(toolsModule.createCommitTools as unknown as { mockReturnValue: (value: unknown) => void }).mockReturnValue([]);
const model = getBundledModel("anthropic", "claude-sonnet-4-5");
if (!model) {
throw new Error("Expected claude-sonnet-4-5 model to exist");
}
await runCommitAgentSession({
cwd: "/tmp",
git: {} as never,
model,
settings: Settings.isolated(),
modelRegistry: {} as never,
authStorage: {} as never,
changelogTargets: [],
requireChangelog: false,
});
expect(prompts).toHaveLength(4);
for (const prompt of prompts) {
expect(prompt.options?.attribution).toBe("agent");
expect(prompt.options?.expandPromptTemplates).toBe(false);
}
});
});
packages/coding-agent/test/model-registry.test.ts
+38
View File
@@ -909,6 +909,44 @@ describe("ModelRegistry", () => {
expect(llama?.reasoning).toBe(false);
});
test("discovers ollama context window from show model_info", async () => {
using _hook = hookFetch((input, init) => {
const url = String(input);
if (url === "http://127.0.0.1:11434/api/tags") {
return new Response(JSON.stringify({ models: [{ name: "gemma3:4b" }] }), {
status: 200,
headers: { "Content-Type": "application/json" },
});
}
if (url === "http://127.0.0.1:11434/api/show") {
const body = JSON.parse(String(init?.body ?? "{}")) as { model?: string };
if (body.model === "gemma3:4b") {
return new Response(
JSON.stringify({
model_info: {
"gemma3.context_length": 131072,
},
}),
{
status: 200,
headers: { "Content-Type": "application/json" },
},
);
}
}
throw new Error(`Unexpected request: ${url}`);
});
const registry = new ModelRegistry(authStorage, modelsJsonPath);
await registry.refresh();
const gemma = registry.find("ollama", "gemma3:4b");
expect(gemma?.contextWindow).toBe(131072);
expect(gemma?.maxTokens).toBe(8192);
expect(gemma?.input).toEqual(["text"]);
expect(gemma?.reasoning).toBe(false);
});
test("discovery failure does not fail model registry refresh", async () => {
writeRawModelsJson({
ollama: {