fix: Address memory leaks in API provider clients and Bedrock handler cache

[kiwina]

PR for API Provider Client Disposal (api_provider_clients) and Bedrock Handler Cache (bedrock_699)

This PR addresses two potential memory leaks:

1. Undisposed SDK clients in various API provider handlers, as identified in leak-report/likely/api_provider_clients.md.
2. Unbounded growth of previousCachePointPlacements in AwsBedrockHandler, as identified in leak-report/likely/bedrock_699.md.

Closes: #4238
Closes: #4240

Description

1. API Provider Clients:
API provider handlers (e.g., AnthropicHandler, OpenAIHandler) create SDK client instances but lacked explicit dispose() methods. The Task class, which holds these handlers, also did not explicitly call for their disposal. This could lead to resource leaks. This patch introduces a dispose() method to the BaseProvider class (which is called reflectively by Task) and ensures Task.abortTask() and Task.dispose() call this.api?.dispose().

2. Bedrock Handler Cache:
The AwsBedrockHandler class uses an instance property previousCachePointPlacements to store cache point information. Entries were added but never removed. This patch overrides the dispose method in AwsBedrockHandler to clear this map and calls super.dispose().

Related Bug Reports

• API Provider Clients: fix: Potentially Undisposed SDK Clients in API Providers (api_provider_clients) #4238
• Bedrock Handler Cache: fix: Unbounded map growth in AwsBedrockHandler (bedrock_699) #4240

Changes

• Modified src/api/providers/base-provider.ts:
  • Added a dispose() method that reflectively attempts to call common cleanup methods on a this.client property.
• Modified src/api/index.ts:
  • Added dispose?: () => void to the ApiHandler interface.
• Modified src/core/task/Task.ts:
  • Added this.api?.dispose?.() to abortTask() and the new dispose() methods.
• Modified src/api/providers/bedrock.ts:
  • Overrode the dispose() method in AwsBedrockHandler to clear this.previousCachePointPlacements and call super.dispose().
• Modified src/api/providers/vscode-lm.ts:
  • Added override keyword to countTokens method.

How to Test

For API Provider Clients:

1. Create and run tasks using various API providers.
2. Abort tasks or simulate scenarios where tasks might be disposed of prematurely.
3. Monitor resource usage (e.g., network connections, memory) to verify that SDK clients are being cleaned up.
4. Check console logs for dispose messages from BaseProvider if client cleanup methods are found.

For Bedrock Handler Cache:

1. Create and run multiple tasks using the AWS Bedrock provider with prompt caching enabled.
2. Ensure each task uses a different conversationId.
3. Abort tasks or simulate Task.dispose().
4. Monitor AwsBedrockHandler or add logging to observe the clearing of previousCachePointPlacements.
5. Verify memory usage does not grow indefinitely with the number of conversations.

[daniel-lxs]

Hey @kiwina, I left some suggestions for this one.

Might be worth actually having a proper dispose implementation for each provider.

Curious to hear your thoughts.

[daniel-lxs]

The reflection-based approach using (this as any).client and typeof checks might be problematic for several reasons:

1. Loss of type safety: TypeScript cannot verify that subclasses have a client property
2. Runtime dependency: Relies on property names and method names that might change
3. Brittle assumptions: Assumes all clients will have one of three disposal methods (close, destroy, or dispose)

Consider a more type-safe approach:

// Option 1: Make dispose abstract and let each provider implement it
export abstract class BaseProvider implements ApiHandler {
    abstract dispose(): void;
}

// Option 2: Define a client contract
export abstract class BaseProvider implements ApiHandler {
    protected abstract getClient(): { dispose?: () => void } | undefined;
    
    public dispose(): void {
        const client = this.getClient();
        client?.dispose?.();
    }
}

This would ensure each provider explicitly handles its own cleanup logic with proper type checking.

[daniel-lxs]

It seems that these are being called more than once.

1. abortTask(true) is called, which internally calls this.api?.dispose?.()
2. Then this.api?.dispose?.() is called again immediately after

This could lead to issues if the disposal method isn't idempotent.

[daniel-lxs]

Several providers in the codebase have client properties but don't override dispose():

• AnthropicHandler has private client: Anthropic
• OpenAiHandler has private client: OpenAI
• BaseOpenAiCompatibleProvider has private client: OpenAI
• Other OpenAI-compatible providers (Groq, Chutes, etc.)

These will rely on this reflection-based disposal, which:

1. Won't work for private properties (they're not accessible via reflection)
2. Assumes the SDK has one of the three 3 methods

As I said before, I think we should probably have each provider implement its own dispose() method to properly clean up its specific resources, otherwise this implementation doesn't actually do much for providers that don't have any of the 3 disposal methods.

[kiwina]

Lets go with a dispose method on the provider, like i mentioned on another pr this would streamline the codebase to use it.

[kiwina]

I'll try to get it done in the next couple of days

[daniel-lxs]

Suggested change

console.log(`[subtasks] disposing task ${this.taskId}.${this.instanceId}`)

[kiwina]

i think generally we have way to little logging to trace problems. I did mention this before on discord.
we should have generally some better logging to aid in development, if we don't know whats going on it's has to fix something.

[daniel-lxs]

This dispose method will directly conflict with your PR #4233.

[kiwina]

Yes, there are several PR that touch into the same concept of the Task class needing to dispose of everything it launches.
I could potentially put them all togheter, i belive keeping them seperated for now will give a better insight. merging them is trivial.

[kiwina]

Hey @kiwina, I left some suggestions for this one.

Might be worth actually having a proper dispose implementation for each provider.

Curious to hear your thoughts.

Technically speking, i agree, practically speaking they all need to be refactored. My fix takes into account a minimally invasive solution

[hannesrudolph]

stale