Skip to content

Feat/pi deferred extension model resolution#1509

Merged
Wirasm merged 6 commits into
coleam00:devfrom
instigator24:feat/pi-deferred-extension-model-resolution
May 15, 2026
Merged

Feat/pi deferred extension model resolution#1509
Wirasm merged 6 commits into
coleam00:devfrom
instigator24:feat/pi-deferred-extension-model-resolution

Conversation

@instigator24

@instigator24 instigator24 commented Apr 30, 2026

Copy link
Copy Markdown
Contributor

Summary

Describe this PR in 2-5 bullets:

  • Problem: Extension providers (e.g. pi-provider-kiro) register their models on the ModelRegistry during session.bindExtensions(), but the current code calls modelRegistry.find() before bindExtensions() and throws immediately when the model isn't found — blocking all extension-registered models.
  • Why it matters: Enables Archon users to access 20 free models (Claude, DeepSeek, Qwen, Kimi, MiniMax, etc.) through the Kiro API via the pi-provider-kiro extension. Without this change, provider: pi + model: kiro/claude-sonnet-4-6 always throws "Pi model not found."
  • What changed: Two-phase model resolution — phase 1 checks the static catalog before bindExtensions(), phase 2 retries after bindExtensions() when extension providers have registered their models. Also switches ModelRegistry.create()ModelRegistry.inMemory() so extensions can call registerProvider() on a mutable registry.
  • What did not change (scope boundary): Built-in providers (google, anthropic, openai, etc.) follow the exact same code path as before — phase 1 resolves them immediately. The ensurePiPackageDirShim() for compiled binary support is preserved. No changes to any other provider, the registry, config, or the DAG executor.

UX Journey

Before

 User                    Archon Pi Provider         Pi SDK
 ────                    ──────────────────         ──────
 sets model:             modelRegistry.find()
 kiro/claude-sonnet-4-6  → returns undefined
                         ✖ throws "Pi model
                           not found"
                         (workflow fails)

After

 User                    Archon Pi Provider         Pi SDK / Extension
 ────                    ──────────────────         ──────────────────
 sets model:             modelRegistry.find()
 kiro/claude-sonnet-4-6  → returns undefined
                         [logs: deferring to
                          extension resolution]
                         createAgentSession()
                          (no model arg)
                         bindExtensions()  ──────▶  pi-provider-kiro calls
                                                    registerProvider("kiro")
                         [modelRegistry.find()      on mutable registry
                          → returns kiro model]
                         session.setModel() ──────▶ switches to kiro model
                         session.prompt()   ──────▶ streams via Kiro API
 sees response ◀──────── yields chunks

Architecture Diagram

Before

 workflow YAML (provider: pi, model: kiro/...)
       │
       ▼
 PiProvider.sendQuery()
       │
       ├─ ModelRegistry.create(authStorage)  ← immutable, reads models.json
       │       │
       │       ▼
       │  registry.find(kiro, claude-sonnet-4-6)
       │       │
       │       ▼
       │  ✖ NOT FOUND → throw Error
       │
       ├─ (never reached) createAgentSession()
       ├─ (never reached) bindExtensions()
       └─ (never reached) session.prompt()

After

 workflow YAML (provider: pi, model: kiro/...)
       │
       ▼
 [~] PiProvider.sendQuery()
       │
       ├─ [~] ModelRegistry.inMemory(authStorage)  ← mutable registry
       │       │
       │       ▼
       │  registry.find(kiro, claude-sonnet-4-6)
       │       │
       │       ▼
       │  NOT FOUND → log info, defer
       │
       ├─ createAgentSession(no model)
       │
       ├─ bindExtensions()
       │       │
       │       ▼
       │  [+] pi-provider-kiro registers models on modelRegistry
       │
       ├─ [+] registry.find(kiro, claude-sonnet-4-6) → FOUND
       ├─ [+] session.setModel(kiroModel)
       │
       └─ session.prompt() → streams response

Connection inventory (list every module-to-module edge, mark changes):

From To Status Notes
PiProvider ModelRegistry modified create()inMemory() for mutable registry
PiProvider createAgentSession modified model arg now conditional (omitted for extension path)
PiProvider session.setModel new deferred model switch after bindExtensions
PiProvider session.bindExtensions unchanged still called in same position
PiProvider AuthStorage unchanged
PiProvider event-bridge unchanged
PiProvider options-translator unchanged
pi-provider-kiro extension ModelRegistry new (runtime) registers provider during bindExtensions

Label Snapshot

  • Risk: risk: low
  • Size: size: M
  • Scope: providers
  • Module: providers:community/pi

Change Metadata

  • Change type: feature
  • Primary scope: providers

Linked Issue

Validation Evidence (required)

Commands and result summary:

bun x tsc --noEmit -p packages/providers/tsconfig.json  # ✅ clean
bun test src/community/pi/provider.test.ts               # ✅ 59 pass, 0 fail
bun test src/registry.test.ts                             # ✅ 32 pass, 0 fail
  • Evidence provided: All 59 Pi provider tests pass. Type-check clean. End-to-end workflow tested with two different Kiro models (see Human Verification below).
  • If any command is intentionally skipped, explain why: bun run validate not run at monorepo root — only the packages/providers package is affected. Full CI will cover the rest.

Security Impact (required)

  • New permissions/capabilities? No
  • New external network calls? No (extension providers make their own calls; this PR doesn't add any)
  • Secrets/tokens handling changed? No (auth fail-fast is now conditional — skipped for extension providers that manage their own credentials — but no new secret handling)
  • File system access scope changed? No

Compatibility / Migration

  • Backward compatible? Yes — built-in providers follow the identical code path. ModelRegistry.inMemory() exposes the same find() interface as create().
  • Config/env changes? No
  • Database migration needed? No

Human Verification (required)

What was personally validated beyond CI:

  • Verified scenarios:
    • archon workflow run archon-test-pi with kiro/claude-sonnet-4-6 — ✅ correct factual answer + verification node passes
    • archon workflow run archon-test-pi with kiro/minimax-m2-5 — ✅ correct factual answer + verification node passes
    • Built-in model google/gemini-2.5-pro — ✅ still works (no regression)
    • Model not found in static catalog AND not found after bindExtensions — ✅ throws clear error with extension install hint
  • Edge cases checked:
    • Session resume on second node (verify node) — no misleading "Could not restore model" warning (suppressed for deferred resolution path)
    • enableExtensions: false — extension models correctly fail with "not found" error
    • models.json load error — warning logged, deferred resolution still attempted
  • What was not verified: ModelRegistry.inMemory() behavior with ~/.pi/agent/models.json custom models (LM Studio, ollama) — no local setup to test, but inMemory() still exposes the static catalog via find() per Pi SDK docs.

Side Effects / Blast Radius (required)

  • Affected subsystems/workflows: Only packages/providers/src/community/pi/provider.ts — the Pi community provider's sendQuery() method.
  • Potential unintended effects: ModelRegistry.inMemory() may not load ~/.pi/agent/models.json custom models the way ModelRegistry.create() does. Users relying on custom models.json entries should verify after upgrade. The static built-in catalog is unaffected.
  • Guardrails/monitoring for early detection: pi.model_not_in_static_catalog_deferring log event fires whenever phase 1 misses — searchable in logs. pi.model_registry_load_error warns on models.json issues.

Rollback Plan (required)

  • Fast rollback command/path: Revert this single commit — change inMemory back to create, remove the if (!model) deferred resolution block, restore hard-fail in step 3.
  • Feature flags or config toggles (if any): enableExtensions: false in ~/.archon/config.yaml disables extension loading entirely, which effectively disables the deferred resolution path.
  • Observable failure symptoms: If rollback is needed, users will see "Pi model not found" errors for extension-registered models (kiro/*). Built-in models are unaffected.

Risks and Mitigations

  • Risk: ModelRegistry.inMemory() may not discover custom models from ~/.pi/agent/models.json that ModelRegistry.create() loaded.

  • Risk: session.setModel() called after createAgentSession() but before prompt() — relies on Pi SDK allowing model switch mid-session.

    • Mitigation: Validated end-to-end with two different Kiro models. Pi SDK's setModel() is a documented API. No prompt has been sent yet when setModel() is called, so the session state is clean.

Summary by CodeRabbit

  • New Features

    • Added example workflows that run parallel Pi and Ollama prompts with a verification step emitting one of four fixed verification statuses.
  • Documentation

    • Added a README with installation and CLI/provider setup for running Pi-based example workflows.
  • Bug Fixes

    • Improved Pi provider model resolution, deferred resolution until extensions bind, reduced false auth failures, and added clearer error logging.
  • Tests

    • Updated tests to reflect revised model-resolution and session behavior.

Review Change Stack

@coderabbitai

coderabbitai Bot commented Apr 30, 2026

Copy link
Copy Markdown

Note

Reviews paused

It looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

  • @coderabbitai resume to resume automatic reviews.
  • @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

  • ▶️ Resume reviews
  • 🔍 Trigger review
📝 Walkthrough

Walkthrough

Implements two-phase model resolution in the Pi provider (static lookup, defer, bindExtensions retry + setModel), adjusts auth timing and session creation, adds error logging in event-bridge, updates provider tests and mocks, and adds Pi workflow examples and a README.

Changes

Examples & Docs

Layer / File(s) Summary
Install README
examples/workflows/README.md
Adds Prerequisites and three bash install snippets for Pi CLI, pi-provider-kiro and required Pi packages, and pi-mcp-adapter.
Workflow examples
examples/workflows/archon-test-pi.yaml, .archon/workflows/archon-test-pi.yaml
Adds archon-test-pi workflow with two parallel prompts (hello via pi, hello-ollama via ollama) and a dependent verify node that emits exactly one of four fixed verdict strings based on whether outputs mention “Detroit Tigers”.

Pi Provider Implementation

Layer / File(s) Summary
Model resolution & session
packages/providers/src/community/pi/provider.ts
Two-phase lookup: call modelRegistry.find() initially and defer (log) if not found; log pi.model_registry_load_error when load/parse errors exist; only perform auth getApiKey when model resolved in static catalog; create session without model when unresolved; suppress fallback system chunk unless model present; after session.bindExtensions() retry find(), dispose+throw with guidance if still missing, and call session.setModel(model) when resolved.

Event Bridge

Layer / File(s) Summary
Result chunk error logging
packages/providers/src/community/pi/event-bridge.ts
buildResultChunk logs an error for terminal error/aborted stops recording stopReason and errorMessage via getLog().error(...).

Pi Provider Tests

Layer / File(s) Summary
Tests & mocks
packages/providers/src/community/pi/provider.test.ts
Tests updated to add/reset mockSetModel, reset mocks per test, and adjust expectations to account for deferred resolution and retry after bindExtensions() (double find() calls and pi.model_registry_load_error warnings).

Sequence Diagram

sequenceDiagram
    participant Client
    participant PiProvider
    participant ModelRegistry
    participant AuthStorage
    participant Session
    participant Extensions

    Client->>PiProvider: initialize(request with model)
    PiProvider->>ModelRegistry: find(model) [phase 1]
    alt model in static catalog
        ModelRegistry-->>PiProvider: model info
        PiProvider->>AuthStorage: getApiKey(model)
        AuthStorage-->>PiProvider: apiKey
    else not found
        ModelRegistry-->>PiProvider: undefined
        PiProvider-->>PiProvider: log deferred resolution
    end

    PiProvider->>Session: createAgentSession(registry, maybe without model)
    Session-->>PiProvider: session

    PiProvider->>Session: bindExtensions()
    Session->>Extensions: load/enable extensions
    Extensions-->>Session: extensions registered

    PiProvider->>ModelRegistry: find(model) [phase 2 retry]
    alt model resolved by extensions
        ModelRegistry-->>PiProvider: model info
        PiProvider->>Session: setModel(model)
        Session-->>PiProvider: model set
    else still unresolved
        PiProvider-->>Client: throw error instructing extension install/enable
    end
Loading

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~45 minutes

Possibly related PRs

  • coleam00/Archon#1284: Overlaps Pi provider and test changes; similar shift to ModelRegistry-based resolution and deferred auth behavior.
  • coleam00/Archon#1559: Related edits on Pi provider query/session path and SettingsManager merging that intersect provider session logic.
  • coleam00/Archon#1355: Related provider rewiring affecting model lookup and dynamic SDK import paths used by the provider.

Poem

🐇 I nudged the registry, then hopped aside,
Extensions woke and helped the models hide,
Deferred a lookup, then bound with cheer,
setModel called — prompts now clear,
Hooray — the workflows hum, carrots near!

🚥 Pre-merge checks | ✅ 5
✅ Passed checks (5 passed)
Check name Status Explanation
Title check ✅ Passed The title 'Feat/pi deferred extension model resolution' accurately and specifically describes the main change: implementing a two-phase model resolution system that defers extension provider model lookup to after session.bindExtensions().
Description check ✅ Passed The PR description is comprehensive and well-structured, covering all major template sections: clear problem statement, UX before/after flows, detailed architecture diagrams, validation evidence with passing tests, security assessment, compatibility confirmation, human verification with end-to-end testing, and documented risks/mitigations.
Docstring Coverage ✅ Passed Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.

✏️ Tip: You can configure your own custom pre-merge checks in the settings.

✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests

Tip

💬 Introducing Slack Agent: The best way for teams to turn conversations into code.

Slack Agent is built on CodeRabbit's deep understanding of your code, so your team can collaborate across the entire SDLC without losing context.

  • Generate code and open pull requests
  • Plan features and break down work
  • Investigate incidents and troubleshoot customer tickets together
  • Automate recurring tasks and respond to alerts with triggers
  • Summarize progress and report instantly

Built for teams:

  • Shared memory across your entire org—no repeating context
  • Per-thread sandboxes to safely plan and execute work
  • Governance built-in—scoped access, auditability, and budget controls

One agent for your entire SDLC. Right inside Slack.

👉 Get started


Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

Comment @coderabbitai help to get the list of available commands and usage tips.

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🧹 Nitpick comments (4)
examples/workflows/README.md (2)

14-19: ⚡ Quick win

Add language specifier to fenced code block.

The code block showing the workflow diagram should have a language specifier (e.g., text or plaintext) for consistency and accessibility.

📝 Suggested fix
-```
+```text
 setup-env → fetch-and-classify → investigate|plan → approve-plan ⏸
 → create-subissues → implement → validate → approve-implementation ⏸
 → create-pr → review-and-classify → [parallel review agents]
 → synthesize → self-fix-and-simplify → linear-finalize
</details>

<details>
<summary>🤖 Prompt for AI Agents</summary>

Verify each finding against the current code and only fix it if needed.

In @examples/workflows/README.md around lines 14 - 19, Add a language specifier
to the fenced code block containing the workflow diagram so it’s
rendered/accessibly labeled (e.g., change the opening triple backticks before
the workflow lines to ```text); update the block that currently contains
"setup-env → fetch-and-classify → investigate|plan → approve-plan ⏸ ..." so the
fenced code block starts with a language token like text or plaintext.


</details>

---

`72-72`: _⚡ Quick win_

**Use descriptive link text instead of "more".**

The link text "more" doesn't convey what users will find at the destination. This impacts accessibility and usability.

<details>
<summary>📝 Suggested fix</summary>

```diff
-Available models: `claude-sonnet-4-6`, `claude-opus-4-7`, `deepseek-3-2`, `minimax-m2-5`, `qwen3-coder-next`, `auto`, and [more](https://github.com/mikeyobrien/pi-provider-kiro#available-models).
+Available models: `claude-sonnet-4-6`, `claude-opus-4-7`, `deepseek-3-2`, `minimax-m2-5`, `qwen3-coder-next`, `auto`, and [see full model list](https://github.com/mikeyobrien/pi-provider-kiro#available-models).
```
</details>

<details>
<summary>🤖 Prompt for AI Agents</summary>

```
Verify each finding against the current code and only fix it if needed.

In `@examples/workflows/README.md` at line 72, Update the anchor text in the
models list line so it is descriptive instead of "more" — replace
`[more](https://github.com/mikeyobrien/pi-provider-kiro#available-models)` with
a meaningful label such as `[available
models](https://github.com/mikeyobrien/pi-provider-kiro#available-models)` or
`[provider available
models](https://github.com/mikeyobrien/pi-provider-kiro#available-models)` in
the README line that lists Available models (`claude-sonnet-4-6`,
`claude-opus-4-7`, `deepseek-3-2`, `minimax-m2-5`, `qwen3-coder-next`, `auto`)
so the link clearly describes its destination for accessibility and usability.
```

</details>

</blockquote></details>
<details>
<summary>packages/providers/src/community/pi/provider.test.ts (1)</summary><blockquote>

`393-407`: _⚡ Quick win_

**Consider adding a test that verifies `session.setModel()` is called for extension-provided models.**

The tests verify the error path when `find()` returns undefined twice, but there's no test covering the happy path where `find()` returns undefined initially (phase 1), then returns a model after `bindExtensions()` (phase 2), triggering `session.setModel()`. This is the primary success path for extension-provided models like `kiro/minimax-m2-5`.

<details>
<summary>💡 Example test structure</summary>

```typescript
test('deferred model resolution calls session.setModel after bindExtensions', async () => {
  process.env.GEMINI_API_KEY = 'sk-test';
  // First find() returns undefined (not in static catalog)
  mockModelRegistryFind.mockImplementationOnce(() => undefined);
  // Second find() after bindExtensions returns the model
  mockModelRegistryFind.mockImplementationOnce((provider, modelId) => ({
    id: modelId,
    provider,
    name: `${provider}/${modelId}`,
  }));
  resetScript(scriptedAgentEnd());

  const { error } = await consume(
    new PiProvider().sendQuery('hi', '/tmp', undefined, {
      model: 'kiro/claude-sonnet-4-6',
    })
  );

  expect(error).toBeUndefined();
  expect(mockSetModel).toHaveBeenCalledTimes(1);
  expect(mockSetModel).toHaveBeenCalledWith(
    expect.objectContaining({ provider: 'kiro' })
  );
});
```
</details>

<details>
<summary>🤖 Prompt for AI Agents</summary>

```
Verify each finding against the current code and only fix it if needed.

In `@packages/providers/src/community/pi/provider.test.ts` around lines 393 - 407,
Add a happy-path unit test that verifies deferred model resolution triggers
session.setModel: when testing PiProvider.sendQuery, mockModelRegistryFind to
first return undefined (static catalog miss) and then return a model object on
the subsequent call (after bindExtensions), call
resetScript(scriptedAgentEnd()), invoke new PiProvider().sendQuery(...) with an
extension-provided model id (e.g., "kiro/..."), then assert no error and that
mockSetModel (or session.setModel) was called exactly once with an object whose
provider equals the extension provider (e.g., 'kiro'); use the existing mocks
(mockModelRegistryFind, mockSetModel) and ensure the second
mockModelRegistryFind implementation returns an object containing id, provider,
and name so session.setModel is exercised.
```

</details>

</blockquote></details>
<details>
<summary>packages/providers/src/community/pi/provider.ts (1)</summary><blockquote>

`262-286`: _💤 Low value_

**Step numbering inconsistency: "4" is reused.**

The comment at line 262 starts with "4. Resolve credentials" but line 326 also starts with "4. Translate Archon nodeConfig". This makes the step sequence confusing. Steps 4e, 4f, and 4g reference a different "step 4" than the credential resolution block.

Consider renumbering to use distinct step numbers (e.g., credentials as step 3, nodeConfig translation as step 4) or converting to descriptive section headers instead of numbered steps.

<details>
<summary>🤖 Prompt for AI Agents</summary>

```
Verify each finding against the current code and only fix it if needed.

In `@packages/providers/src/community/pi/provider.ts` around lines 262 - 286, The
comment block labeled "4. Resolve credentials" duplicates a "step 4" used later
for "4. Translate Archon nodeConfig", causing numbering confusion; update the
comment headings to unique identifiers (e.g., renumber credentials to "3.
Resolve credentials" or replace numeric steps with descriptive headers) and
adjust any sub-step labels (4e/4f/4g) to match the new scheme so references
remain consistent; ensure you update the surrounding comments that reference
parsed.provider, PI_PROVIDER_ENV_VARS, envOverride, and
authStorage.setRuntimeApiKey to reflect the new step/header naming.
```

</details>

</blockquote></details>

</blockquote></details>

<details>
<summary>🤖 Prompt for all review comments with AI agents</summary>

Verify each finding against the current code and only fix it if needed.

Nitpick comments:
In @examples/workflows/README.md:

  • Around line 14-19: Add a language specifier to the fenced code block
    containing the workflow diagram so it’s rendered/accessibly labeled (e.g.,
    change the opening triple backticks before the workflow lines to ```text);
    update the block that currently contains "setup-env → fetch-and-classify →
    investigate|plan → approve-plan ⏸ ..." so the fenced code block starts with a
    language token like text or plaintext.
  • Line 72: Update the anchor text in the models list line so it is descriptive
    instead of "more" — replace
    [more](https://github.com/mikeyobrien/pi-provider-kiro#available-models) with
    a meaningful label such as [available models](https://github.com/mikeyobrien/pi-provider-kiro#available-models) or
    [provider available models](https://github.com/mikeyobrien/pi-provider-kiro#available-models) in
    the README line that lists Available models (claude-sonnet-4-6,
    claude-opus-4-7, deepseek-3-2, minimax-m2-5, qwen3-coder-next, auto)
    so the link clearly describes its destination for accessibility and usability.

In @packages/providers/src/community/pi/provider.test.ts:

  • Around line 393-407: Add a happy-path unit test that verifies deferred model
    resolution triggers session.setModel: when testing PiProvider.sendQuery,
    mockModelRegistryFind to first return undefined (static catalog miss) and then
    return a model object on the subsequent call (after bindExtensions), call
    resetScript(scriptedAgentEnd()), invoke new PiProvider().sendQuery(...) with an
    extension-provided model id (e.g., "kiro/..."), then assert no error and that
    mockSetModel (or session.setModel) was called exactly once with an object whose
    provider equals the extension provider (e.g., 'kiro'); use the existing mocks
    (mockModelRegistryFind, mockSetModel) and ensure the second
    mockModelRegistryFind implementation returns an object containing id, provider,
    and name so session.setModel is exercised.

In @packages/providers/src/community/pi/provider.ts:

  • Around line 262-286: The comment block labeled "4. Resolve credentials"
    duplicates a "step 4" used later for "4. Translate Archon nodeConfig", causing
    numbering confusion; update the comment headings to unique identifiers (e.g.,
    renumber credentials to "3. Resolve credentials" or replace numeric steps with
    descriptive headers) and adjust any sub-step labels (4e/4f/4g) to match the new
    scheme so references remain consistent; ensure you update the surrounding
    comments that reference parsed.provider, PI_PROVIDER_ENV_VARS, envOverride, and
    authStorage.setRuntimeApiKey to reflect the new step/header naming.

</details>

---

<details>
<summary>ℹ️ Review info</summary>

<details>
<summary>⚙️ Run configuration</summary>

**Configuration used**: defaults

**Review profile**: CHILL

**Plan**: Pro

**Run ID**: `8258c687-d2d9-4b03-8c25-a37e726cfdcb`

</details>

<details>
<summary>📥 Commits</summary>

Reviewing files that changed from the base of the PR and between 8295ece76311a55264ee4c8e9af2c8395b626e50 and 1fbdee30e9835079ae9d2c3dd59c9de1a12ab16a.

</details>

<details>
<summary>📒 Files selected for processing (4)</summary>

* `examples/workflows/README.md`
* `examples/workflows/archon-test-pi.yaml`
* `packages/providers/src/community/pi/provider.test.ts`
* `packages/providers/src/community/pi/provider.ts`

</details>

</details>

<!-- This is an auto-generated comment by CodeRabbit for review status -->

@instigator24 instigator24 force-pushed the feat/pi-deferred-extension-model-resolution branch from 1fbdee3 to ab5c9bf Compare April 30, 2026 22:18

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🧹 Nitpick comments (1)
packages/providers/src/community/pi/provider.ts (1)

475-476: ⚡ Quick win

Update stale mutability comment for ModelRegistry.create().

These lines say create() is immutable and inMemory() is required, but this file now intentionally relies on ModelRegistry.create() being mutable (also documented earlier around Line 207-Line 212). Please align this comment to prevent future regressions.

Proposed comment fix
-    //     ModelRegistry.inMemory() is required (step 2) — create() returns
-    //     an immutable registry that registerProvider() can't modify.
+    //     This works with ModelRegistry.create() (step 2): extension
+    //     session_start hooks can call registerProvider() and mutate the
+    //     same registry instance used by this provider.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@packages/providers/src/community/pi/provider.ts` around lines 475 - 476, The
comment claiming ModelRegistry.create() returns an immutable registry (and that
ModelRegistry.inMemory() is required) is stale; update the comment to state that
ModelRegistry.create() produces a mutable registry that registerProvider() can
modify (consistent with the earlier ModelRegistry.create() documentation) and
remove or clarify the note about inMemory() being required so future readers
won't revert this intentional behavior; locate the comment near the existing
references to ModelRegistry.create(), ModelRegistry.inMemory(), and
registerProvider() and rewrite it to accurately describe mutability and intended
usage.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.

Nitpick comments:
In `@packages/providers/src/community/pi/provider.ts`:
- Around line 475-476: The comment claiming ModelRegistry.create() returns an
immutable registry (and that ModelRegistry.inMemory() is required) is stale;
update the comment to state that ModelRegistry.create() produces a mutable
registry that registerProvider() can modify (consistent with the earlier
ModelRegistry.create() documentation) and remove or clarify the note about
inMemory() being required so future readers won't revert this intentional
behavior; locate the comment near the existing references to
ModelRegistry.create(), ModelRegistry.inMemory(), and registerProvider() and
rewrite it to accurately describe mutability and intended usage.

ℹ️ Review info
⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: 45018f4b-e6a8-4db2-8291-b31e3bb9c9aa

📥 Commits

Reviewing files that changed from the base of the PR and between ab5c9bf and fdeb783.

📒 Files selected for processing (4)
  • .archon/workflows/archon-test-pi.yaml
  • examples/workflows/archon-test-pi.yaml
  • packages/providers/src/community/pi/event-bridge.ts
  • packages/providers/src/community/pi/provider.ts
✅ Files skipped from review due to trivial changes (1)
  • .archon/workflows/archon-test-pi.yaml
🚧 Files skipped from review as they are similar to previous changes (1)
  • examples/workflows/archon-test-pi.yaml

…esolution

Extension providers (e.g. pi-provider-kiro) register their models on
the ModelRegistry during session.bindExtensions(), not during the
initial modelRegistry.find() call. The previous approach threw
immediately when find() returned undefined, blocking extension models.

New flow:
1. modelRegistry.find() checks static catalog + models.json (step 3)
2. If found: pass to createAgentSession, fail-fast auth as before
3. If not found: log info, skip auth (extension providers manage their
   own credentials), create session without model
4. After bindExtensions() (step 4g): retry modelRegistry.find() — now
   extension-registered models are discoverable — and call
   session.setModel() to switch to the resolved model
5. If still not found after bindExtensions(): throw with a hint about
   installing the provider extension with enableExtensions: true

This preserves all existing behavior for built-in providers and
models.json custom models while adding support for extension-registered
providers like pi-provider-kiro (Kiro API — 20 free models including
Claude, DeepSeek, Qwen, etc.).

Changes:
- provider.ts: deferred model resolution after bindExtensions();
  conditional model arg to createAgentSession; conditional auth
  fail-fast; session.setModel() for extension-resolved models
- provider.test.ts: add setModel mock; update model-not-found tests
  for two-phase find() behavior
Simple two-node workflow that verifies the Pi provider works with
extension-registered models (e.g. pi-provider-kiro). Node 1 asks a
factual question, node 2 validates the answer.

Tested with kiro/claude-sonnet-4-6 and kiro/minimax-m2-5.
…ages

- Switch ModelRegistry.inMemory() to ModelRegistry.create() so custom
  providers from ~/.pi/agent/models.json (ollama, LM Studio, etc.)
  resolve correctly. Both return the same mutable class supporting
  registerProvider() for extensions.
- Surface Pi's AssistantMessage.errorMessage in the result chunk errors
  array so dag-executor shows actual errors instead of generic
  'SDK returned error'.
- Add ollama/qwen3.5 node to archon-test-pi workflow to validate both
  extension (kiro) and models.json (ollama) provider paths.
@instigator24 instigator24 force-pushed the feat/pi-deferred-extension-model-resolution branch from fdeb783 to a06443c Compare May 9, 2026 14:24
@Wirasm

Wirasm commented May 14, 2026

Copy link
Copy Markdown
Collaborator

Review Summary

Verdict: minor-fixes-needed

This PR implements two-phase model resolution for extension-registered Pi providers, enabling kiro/* and similar extension models to work correctly. The core logic is clean, all 59 tests pass, and the error messages include actionable hints. Two items need attention before merge: an inconsistency in event-bridge.ts where error logging returns the chunk without signaling failure, and several verbose multi-paragraph comments that document internal step numbers in a way that will drift when the code is refactored.

Blocking issues

  • event-bridge.ts:160–164: The isError branch logs at error level but returns the chunk the same way as the non-error path. Callers have no way to distinguish a successfully-processed chunk from one that represents an error condition — the result still flows into the message stream. Fix: add an explicit comment documenting this as intentional design, or propagate an error indicator in the chunk type field.

Suggested fixes

  • provider.ts:299–304: The auth fail-fast is gated on if (model) with no comment explaining why extension providers skip this check. A future developer could add a new built-in provider and unknowingly hit the deferred path, skipping credential validation entirely. Fix: add a comment at the if (model) guard explaining which providers are in the static catalog vs. registered by extensions, or add a isStaticCatalogModel boolean to make the branching explicit.

  • provider.ts (multiple locations — lines 239–259, 268–284, 325–351, 551–562, 568–586): Five multi-paragraph comments document internal step numbers ("step 2", "step 3", "step 4f", "step 4g") and Pi SDK internals (providerActions, ExtensionRunner, session restore quirks). These will drift when the code is refactored. The phrase "Despite the earlier assumption" documents history, not a design invariant. Fix: replace each block with one-sentence WHY explanations:

    • 239–259: "Extension providers need an in-memory registry so they can call registerProvider() during bindExtensions()."
    • 268–284: "[LOOKUP-1] Check the static catalog first (phase 1 of 2)."
    • 325–351: "Auth validation deferred for extension providers — they manage credentials outside Pi's AuthStorage."
    • 551–562: "Extension providers register their models during bindExtensions() — this is why we pass the in-memory registry in step 2."
    • 568–586: "[LOOKUP-2] Re-check the registry after bindExtensions() for extension-registered models. Safe because no prompt has been sent yet."
  • provider.ts:530–534: The modelFallbackMessage suppression guard lacks a comment. Future readers won't know that extension models aren't in the static catalog and that Pi's session restore will misleadingly warn about that. Fix: "Extension models aren't in the static catalog — skip the fallback warning."

  • provider.test.ts:332–333 and 411–412: Test comments reference step numbers ("step 3 and step 4g") that will drift. Fix: "find() is called twice — first on the static catalog, then again after bindExtensions() resolves extension providers."

Minor / nice-to-have

  • event-bridge.ts event naming: pi.event-bridge.result_error uses a period instead of underscore, breaking the {domain}.{action}_{state} convention. Rename to pi.event_bridge.result_error or pi.result_chunk_error.

  • event-bridge.ts:163: model: last.model in the error log could be noisy since models can contain arbitrary strings. Consider omitting it or documenting the intent.

  • provider.ts:567–592: Session cleanup gap — if session.setModel(model) throws after the session is created, there's no explicit session.dispose(). Unlikely to matter in practice, but a try-catch or comment would be cleaner.

  • provider.ts:287: Comment typo — extra space before "class" reads awkwardly. Fix: create() returns the same ModelRegistry class as inMemory() —

Compliments

  • The deferred resolution pattern (model omitted from createAgentSession, set later via session.setModel()) is well-documented in a single concise comment at lines 520–522.
  • The PR description's UX journey and architecture diagrams are excellent — rare to see this level of clarity in a PR.
  • provider.test.ts:339 ("Both must return undefined to trigger the error") explains a non-obvious test requirement well.

Reviewed via maintainer-review-pr workflow (Pi/Minimax). Aspects run: code-review, error-handling, test-coverage, comment-quality.

… test mock alignment

- event-bridge: document intentional design of yielding error chunks (isError:true
  propagates in chunk, not via throw); fix event name to pi.result_chunk_error;
  remove noisy model field from error log
- provider: replace verbose multi-paragraph step-numbered comments with concise
  WHY explanations; fix stale ModelRegistry.inMemory() reference (code uses create());
  add session.dispose() on model-not-found throw and setModel() failure to close
  the session cleanup gap
- provider.test: align mock from ModelRegistry.inMemory → ModelRegistry.create to
  match actual provider code (was already broken on the branch); update test
  description and comments to describe behavior rather than step numbers

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
@jhegeman-ds

Copy link
Copy Markdown
Contributor

@Wirasm — thanks for the thorough review. Here's a summary of everything addressed in commit 979ae7a:

Blocking

event-bridge.ts:160–164 — isError chunk return path
Added an explicit "Intentional design" comment: error chunks are yielded (not thrown) because isError: true in the chunk is the signal. Callers (bridgeSession, dag-executor) check result.isError to classify failures and still receive full token/stopReason context from the same chunk.


Suggested fixes

provider.ts — verbose multi-paragraph step-numbered comments (5 locations)
Replaced all five blocks with concise single-line WHY explanations:

  • Step 2 (AuthStorage/ModelRegistry): clarifies create() is mutable and why extension providers can call registerProvider() during bindExtensions()
  • Step 3 (LOOKUP-1): labeled as [LOOKUP-1] with a one-liner deferring to [LOOKUP-2] for extension providers
  • Step 4 credentials block: trimmed to the essential merge behavior (request env vars win over auth.json)
  • Auth fail-fast comment: replaced 9-line block with a 3-line explanation of why extension providers skip Pi's AuthStorage check
  • Step 4f bindExtensions: trimmed 11 lines (including the stale ModelRegistry.inMemory() reference) to 2 lines
  • Step 4g (LOOKUP-2): labeled to match LOOKUP-1, notes safety of calling setModel() before first prompt

provider.ts:530–534 — modelFallbackMessage suppression
Replaced 4-line comment with: // Extension models aren't in the static catalog — skip the fallback warning.

provider.test.ts:332–333 and 411–412 — step number references
Replaced both with: // find() is called twice — first on the static catalog, then again after bindExtensions() resolves extension providers.


Minor / nice-to-have

event-bridge.ts event naming
Renamed pi.event-bridge.result_errorpi.result_chunk_error to follow the {domain}.{action}_{state} convention.

event-bridge.ts:163 — noisy model field in error log
Removed model: last.model from the error log payload.

provider.ts:567–592 — session cleanup gap
Added session.dispose() on the model-not-found throw, and a try/catch around session.setModel() that disposes the session if setModel() throws — closes the gap where a session could be created but never cleaned up.

provider.ts:287 — comment typo
The comment containing the extra space was inside one of the verbose blocks that was replaced entirely — resolved by deletion.


Test mock alignment (pre-existing breakage)

The last commit (a06443c0) switched the provider from ModelRegistry.inMemory() to ModelRegistry.create() but didn't update the test mocks — the tests were already broken on the branch (mock stub never fired, causing create is not a function errors across 50 tests). Updated mockModelRegistryInMemorymockModelRegistryCreate throughout and corrected the mock module registration. All 65 provider tests pass.

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@packages/providers/src/community/pi/provider.test.ts`:
- Line 46: Add a test exercising the deferred-resolution success path: configure
mockModelRegistryFind to return undefined on first call and a model object on
the second, then call new PiProvider().sendQuery(...) (using
resetScript(scriptedAgentEnd()) and the same options pattern as other tests) and
assert that mockSetModel (the session.setModel mock) was called once with an
object containing id: 'custom-model' and provider: 'extension-provider', and
that mockModelRegistryFind was called twice and the returned error is undefined;
this verifies ModelRegistry.find → undefined then model after bindExtensions
leads to session.setModel being invoked.
🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

  • Push a commit to this branch (recommended)
  • Create a new PR with the fixes

ℹ️ Review info
⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: 789bf758-6ccd-457f-af57-f01630eca528

📥 Commits

Reviewing files that changed from the base of the PR and between a06443c and 979ae7a.

📒 Files selected for processing (3)
  • packages/providers/src/community/pi/event-bridge.ts
  • packages/providers/src/community/pi/provider.test.ts
  • packages/providers/src/community/pi/provider.ts
🚧 Files skipped from review as they are similar to previous changes (2)
  • packages/providers/src/community/pi/event-bridge.ts
  • packages/providers/src/community/pi/provider.ts

Comment thread packages/providers/src/community/pi/provider.test.ts
Verifies that when find() returns undefined on the first call (static catalog miss)
and a model on the second call (after bindExtensions()), session.setModel() is
invoked with the resolved model and no error is returned.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
@jhegeman-ds

Copy link
Copy Markdown
Contributor

@Wirasm — one more follow-up in commit 9081bdd:

provider.test.ts — added happy-path test for deferred extension model resolution

Added the test that was flagged by both your review and CodeRabbit: verifies the full success path where find() returns undefined on the first call (static catalog miss) and a model object on the second call (after bindExtensions()). Asserts that:

  • mockModelRegistryFind is called exactly twice
  • session.setModel() is called once with the resolved model ({ id: 'custom-model', provider: 'extension-provider' })
  • No error is returned

All 66 provider tests pass.

@instigator24

Copy link
Copy Markdown
Contributor Author

(Reposting from the correct account — the previous comment from jhegeman-ds was also us.)

@Wirasm — one more follow-up in commit 9081bdd:

provider.test.ts — added happy-path test for deferred extension model resolution

Added the test that was flagged by both your review and CodeRabbit: verifies the full success path where find() returns undefined on the first call (static catalog miss) and a model object on the second call (after bindExtensions()). Asserts that:

  • mockModelRegistryFind is called exactly twice
  • session.setModel() is called once with the resolved model ({ id: 'custom-model', provider: 'extension-provider' })
  • No error is returned

All 66 provider tests pass.

@Wirasm Wirasm merged commit 7bdf931 into coleam00:dev May 15, 2026
1 check passed
cropse pushed a commit to cropse/Archon that referenced this pull request May 19, 2026
* feat(pi): support extension-registered provider models via deferred resolution

Extension providers (e.g. pi-provider-kiro) register their models on
the ModelRegistry during session.bindExtensions(), not during the
initial modelRegistry.find() call. The previous approach threw
immediately when find() returned undefined, blocking extension models.

New flow:
1. modelRegistry.find() checks static catalog + models.json (step 3)
2. If found: pass to createAgentSession, fail-fast auth as before
3. If not found: log info, skip auth (extension providers manage their
   own credentials), create session without model
4. After bindExtensions() (step 4g): retry modelRegistry.find() — now
   extension-registered models are discoverable — and call
   session.setModel() to switch to the resolved model
5. If still not found after bindExtensions(): throw with a hint about
   installing the provider extension with enableExtensions: true

This preserves all existing behavior for built-in providers and
models.json custom models while adding support for extension-registered
providers like pi-provider-kiro (Kiro API — 20 free models including
Claude, DeepSeek, Qwen, etc.).

Changes:
- provider.ts: deferred model resolution after bindExtensions();
  conditional model arg to createAgentSession; conditional auth
  fail-fast; session.setModel() for extension-resolved models
- provider.test.ts: add setModel mock; update model-not-found tests
  for two-phase find() behavior

* test: add archon-test-pi example workflow for Kiro extension validation

Simple two-node workflow that verifies the Pi provider works with
extension-registered models (e.g. pi-provider-kiro). Node 1 asks a
factual question, node 2 validates the answer.

Tested with kiro/claude-sonnet-4-6 and kiro/minimax-m2-5.

* docs: add Pi/Kiro prerequisites and setup to examples README

* fix(pi): load models.json custom providers and surface SDK error messages

- Switch ModelRegistry.inMemory() to ModelRegistry.create() so custom
  providers from ~/.pi/agent/models.json (ollama, LM Studio, etc.)
  resolve correctly. Both return the same mutable class supporting
  registerProvider() for extensions.
- Surface Pi's AssistantMessage.errorMessage in the result chunk errors
  array so dag-executor shows actual errors instead of generic
  'SDK returned error'.
- Add ollama/qwen3.5 node to archon-test-pi workflow to validate both
  extension (kiro) and models.json (ollama) provider paths.

* fix(pi): address Wirasm PR review — comment clarity, session cleanup, test mock alignment

- event-bridge: document intentional design of yielding error chunks (isError:true
  propagates in chunk, not via throw); fix event name to pi.result_chunk_error;
  remove noisy model field from error log
- provider: replace verbose multi-paragraph step-numbered comments with concise
  WHY explanations; fix stale ModelRegistry.inMemory() reference (code uses create());
  add session.dispose() on model-not-found throw and setModel() failure to close
  the session cleanup gap
- provider.test: align mock from ModelRegistry.inMemory → ModelRegistry.create to
  match actual provider code (was already broken on the branch); update test
  description and comments to describe behavior rather than step numbers

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

* test(pi): add happy-path test for deferred extension model resolution

Verifies that when find() returns undefined on the first call (static catalog miss)
and a model on the second call (after bindExtensions()), session.setModel() is
invoked with the resolved model and no error is returned.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

---------

Co-authored-by: jhegeman-ds <joel.hegeman@deepseas.com>
Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>
@Wirasm Wirasm mentioned this pull request May 28, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants