Hi <@U0AAGJ8MM1C>! Sorry for the late reply. Mayb...

# koog-agentic-framework

[Sergei Pestrikov]

Hi @Chrylo! Sorry for the late reply.

> Maybe the StorageKey should be a data class?
AIAgentStorageKey used to be a data class, but it became an ordinary class last September. The intention to make it truly unique was to prevent accidental storage overwrite in different subgraphs or nodes (e.g. when unrelated subgraphs use the same key named "data"). And also to prevent ClassCastException in the following example if AIAgentStorageKey were a data class:

val stringKey: AIAgentStorageKey<String> = createStorageKey("my_key")
val intKey: AIAgentStorageKey<Int> = createStorageKey("my_key")

// If data class:
storage.set(stringKey, "hello")       // stores "hello"
val value: Int? = storage.get(intKey) // returns "hello" cast to Int? → ClassCastException at use site

> In our case it is cumbersome because we have all the subgraphs in different files and the storage keys are created dynamically.
For this use case it's indeed cumbersome. One solution might be to introduce a key registry (cached factory) to manage shared keys between subgraphs from different files. Let me know if it's not an option (I've added an example of the registry AIAgentStorageKeys at the end of this commit).


> However when I debug into it I see different instances of the "context" for every subgraph so the storage is also different. Is this intended?
True, AIAgentSubgraph.execute() creates a copy of the context via context.replace(context.copy(...)). However, context.storage points to the same storage instance in the copy. I've added some tests in a separate branch to demonstrate that.


> Should I always use the parentcontext? Should we use the root context in this case even though it is internal?
parentcontext and root context are not needed for this case. Copied contexts should share the same underlying storage instance.


One more thing: the documentation about AIAgentStorageKey states that "Each key has a name property which is a string identifier that uniquely represents the storage key." That's misleading, sorry about that, the name property in the current implementation doesn't affect key uniqueness. We'll fix the documentation.

[Chrylo]

Hi @Sergei Pestrikov, I understand the issue with the StorageKey. Even though I think it should be a data class (again). A storage key with the same "key" name IMHO implies always that it will be handled the same by the framework. I personally think a Exception is correct because the same key should not be used for multiple data types, so I would argue that behavior as intended. My take is that in every project, which is a bit more complex, you will outsource the subgraphs quite often. And expecting a lot of devs to write a custom key registry seems not like a good tradeoff.

The main issue for us is: We are already using a cached registry BUT we are developing a REST service with koog. So we are restoring conversations, llm messages and stored node outputs for multiple users. If the StorageKey is not a data class this means we can't restore saved objects when chatting with the agent (or the service is restarted) because the key is lost.

So my main question is: How do I restore storage objects when the key is lost because I can't serialize it.

Thanks again for the good support. :)

[Chrylo]

I also think this issue is related: https://kotlinlang.slack.com/archives/C08SLB97W23/p1772183922623099

[Sergei Pestrikov]

Hi @Chrylo! I'll discuss with the team if we can introduce some sort of registry or to make the StorageKey a data class once again 🙂
Thank you for the details!

About

So we are restoring conversations, llm messages and stored node outputs for multiple users.

A new feature called "ChatMemory" has been merged yesterday into the develop branch. It will allow to run persistent/stateful agents by "agent.run(input, sessionId)" where "sessionId" is a string identifier of a conversation or a user. Let us know if the feature can be useful for your use cases.

[Vadim Briliantov]

And it was also previously possible with Persistence feature + specific flag that tells to only persist messages